Supervised Pretraining for Molecular Force Fields and Properties Prediction
Abstract
Machine learning approaches have become popular for molecular modeling tasks, including molecular force fields and properties prediction.
Traditional supervised learning methods suffer from scarcity of labeled data for particular tasks, motivating the use of large-scale dataset for other relevant tasks.
We propose to pretrain neural networks on a dataset of 86 millions of molecules with atom charges and 3D geometries as inputs and molecular energies as labels.
Experiments show that, compared to training from scratch, fine-tuning the pretrained model can significantly improve the performance for seven molecular property prediction tasks and two force field tasks.
We also demonstrate that the learned representations from the pretrained model contain adequate information about molecular structures, by showing that linear probing of the representations can predict many molecular information including atom types, interatomic distances, class of molecular scaffolds, and existence of molecular fragments.
Our results show that supervised pretraining is a promising research direction in molecular modeling
Introduction
Molecular force fields and properties prediction are important tasks for biotechnology, drug discovery, material and energy science.
Machine learning has become a promising approach for these tasks.
A weakness of supervised machine learning approaches is that they often require a large amount of labeled data to perform well.
However, the molecular datasets are usually small due to the high cost to obtain the labels, which may require wet lab experiments or expensive quantum mechanics simulation
Pretraining is one approach to alleviate the low-data issue by leveraging relevant large-scale data.
The assumption is that the knowledge learned from pretraining can be transferred to the fine-tuning stage on downstream tasks.
For computer vision, the most popular pretraining strategy is probably a supervised learning approach where models are pretrained on image classification tasks using large-scale labeled datasets such as ImageNet [1] and CIFAR [2].
For natural language processing, self-supervised pretraining becomes dominant with the success of the masked language modeling [3] and generative pretraining [4].
Recently, pretraining strategies are explored for molecular modeling [5, 6, 7, 8], and most works employ a self-supervised learning pretraining approach.
Can supervised pretraining help molecular modeling?
Besides the data availability issue, another challenge is to choose a proper supervised learning target.
Hu et al.
[5] experimented with a supervised pretraining strategy, multi-task prediction of more than one thousands biochemical assays.
However they observed negative transfer for a few downstream tasks, and suspected that this is due to such supervised pretraining is not "truly-related" to the downstream tasks.
We argue that the supervised pretraining tasks should focus on more fundamental physical properties, instead of the biochemical properties that are measured in a complex environment and are relevant to only a narrow range of downstream tasks.
We propose to pretrain models to predict molecular energy from molecular structure, with the following considerations.
(i)
The supervised pretraining data should be rich.
Data for fundamental physical properties, such as molecular energies, are often more abundant.
For example, the dataset we use in this work, PubChem PM6 dataset [9], contains energies calculated for 221 millions molecules. (ii)
The supervised pretraining labels should be accurate rather than noisy.
Molecular energy can be calculated using established first principle based approach such as density functional theory [10] or semi-empirical method [11].
Therefore the labels can be obtained by the same mechanism for all samples.
In contrast the experimentally
measured biochemical
assays used in previous work [5] can be noisy due to the difference in experimental conditions or methods across various data sources.
(iii)
The supervised pretraining target should be relevant various downstream tasks.
Many molecular properties of interest are quantities describing the interaction between molecules and the environment, such as proteins in human body or catalyst in batteries.
These processes are governed by the forces acting on the atoms.
Molecular energy is relevant as the force acting on an atom is the negative partial gradients of the molecular potential energy with respect to the coordinate of the atom.
This
close relation between molecular energy and molecular structure encourage the pretrained model to
understand and represent molecular structures, which
is important to perform well on various downstream tasks [12].
We empirically find that the pretrained model learns
to embed various molecular structural information (
see Section 4.3).
We further propose a force regularization technique, which minimizes the magnitude of the energy gradient with respect to atom coordinates.
This makes the pretrained model more suitable for force field prediction tasks.
In contrast, instead of using the exact atom coordinates, most previous pretraining tasks [5, 6, 7, 8] are designed to only use the topological molecular structures.
This partially explain why pretraininig approach was not previously used for force field tasks.
framework.
As illustrated in Figure 1, we find that the proposed pretraining strategy can improve the model performance on various downstream tasks, including molecular force fields and molecular properties prediction such as toxicity and molecular water solubility.
Linear probing shows the learned representation can predict the input atom types, interatomic distance, molecular scaffolds, and the existence of 85 functional groups.
Our contribution is three-folds.
We propose a supervised pretraining strategy which can improve molecular modeling performance.
The pretraining task is to predict
molecular energy , a fundamental physical quantity relevant to various downstream tasks and available in large-scale
dataset
We show that, without manual design or constraints, the pretrained model learn to embed the molecular structural information
This is necessary to a wide range of downstream molecular properties prediction tasks.
We show that the proposed pretraining approach helps force fields models generalize better to unseen molecules
compared
to training
from scratch  Background
Molecular force field.
This task of predicting the force vector acting on each atom in a molecule.
The force depends on the positions of the atoms (conformation).
Chmiela et al.
[13] proposed a MD dataset containing the energy-conservative force fields data for a few species.
An gradient-domain machine learning (GDML) approach is proposed in [13].
More recently, Klicpera et al.
[14] proposed a GemNet model and obtained better results compared to a few previous works.
Pretraining has not been employed in existing works for this task.
Molecular properties There are enormous molecular properties of interest.
Moleculenet [15] collected various datasets of molecular properties on quantum mechanics, physical chemistry [16, 17, 18], biophysics [19], and physiology [20, 21, 22].
Existing works mostly represent the molecule as graphs, with atoms as node and chemical bonds as edges, and then model the molecules with graph neural networks [23, 24, 25, 26] or Transformer[27] based architectures [28, 7].
Molecular substructures Predicting the molecular properties from their structures have been explored extensively [12].
Classic methods often investigate molecule structures from two levels, the scaffold [29], and fragments (or motifs), and it is believed that the molecular properties are significantly affected by these substructures [12].
The scaffold has been used to group the molecules and split the train/validation/test sets to mimic a more realistic and challenging setting [15].
The prediction of the existence of fragments have been used in recent deep learning based methods [7] in a un-supervised approach.
The pretraining-then-fine-tuning paradigm has been used in molecular modeling, with a focus on self-supervised learning.
Hu et al.
[5] experimented with both node-level and graph-level pre-training.
The node-level tasks include prediction of the context of a node, and the prediction of masked node attributes.
For graph-level tasks, they experimented with supervised graphlevel property prediction, and graph structural similarity prediction.
Rong et al.
[7] demonstrated two self-supervised learning pretraining tasks: contextual property prediction and graph-level motif prediction.
Li et al.
[6] proposed a self-supervised pretraining task named pairwise subgraph discrimination, which compares two subgraphs and discriminate whether they come from the same source.
Liu et al.
[8] aligns the latent space learned for molecular graphs and 3D Geometry in attempt to learn knowledge from both input formats.
Method
We consider two families of downstream tasks: molecular force field prediction, and molecular properties prediction.
For the former task, we consider models that take in molecular structure, usually the charges Z (e.g., hydrogen, carbon or oxygen) and 3D coordinates of the atoms R , as input, and output the forces
F acting on each atom.
The force for an atom is a 3D vector, usually sharing the same frame of reference as the atom coordinates.
For the latter task, we consider models that take in molecular structure, and output y , which is a scalar value (for regression tasks) or a probability distribution (for classification tasks).
Our approach follows the popular pretraining-and-finetuning paradigm.
3.1 Pretraining
We use the PubChem PM6 dataset [9] for pretraining.
Energy for 86 millions of optimized molecular 3D geometry at neutral states are provided in this dataset.
One limitation of our approach is that the performance is affected by the choice of pretraining dataset.
The PM6 dataset contain relatively small molecules which makes the pretrained model not suitable for large molecules such as protein or catalysts.
3.1.1 Molecular force fields
We train the model to predict the energy E of the optimized 3D molecular geometry R .
We define a loss term L E based on a distance measure, d , between the predicted energy ˆ E and the energy label value E .
[Equation]
We use mean absolute error as d for L E .
The PubChem PM6 dataset is built for optimized molecules geometry.
The molecular energy is minimized with respect to the atom coordinates.
Therefore the energy gradient with respect to the atom coordinates should be close to zero.
This motivates us to include a regularization loss term
[Equation]
As this gradient is actually negative potential forces acting on the atoms, we refer L ∂ ˆ E/∂R as force regularization.
The pretraining loss for force fields (FF) task is a linear combination of two loss items.
[Equation]
where α is a hyperparameter.
3.1.2 Molecular properties
For the molecular properties prediction tasks, many downstream datasets often provide molecular structures in SMILES formats without exact atom coordinates.
We instead use an estimated 3D geometry R noisy obtained from SMILES using tools such as RDKit [30].
[Equation]
In this case the atom coordinates are approximated, so we do not apply the force regularization term.
The pretraining loss is
[Equation]
3.2 Fine-tuning
During fine-tuning, we use the parameters of the pretrained model to initialize the model parameters, except these for the final output layers.
For the final output layers parameters, we use random initialization.
For molecular properties prediction tasks, the training loss measure the distance d between the label output y and predicted output y .
[Equation]
where the choice of function
d depends on the tasks.
We use mean absolute error for regression tasks, and cross entropy loss for classification tasks.
For molecular force field prediction tasks, we use a linear combination of the energy and force loss.
[Equation]
where γ is a hyperparameter.
Experiments and discussion 4.1 Molecular force fields
For molecular force fields prediction tasks, the existing works [14, 31, 32] generally employ an in-domain setting, where the molecule type in training set is the same as the test set.
Besides this conventional setting, we consider a more challenging, out-of-domain setting.
The models are tested on the force field of a unseen molecule.
This is motivated by distinct characteristics of the pretraining and fine-tuning datasets.
The PubChem PM6 dataset used in pretraining contains a large number of molecules but each only has one conformation.
In contrast, the molecular force field dataset contains many conformations for a limited number of molecules.
We wonder if we can combine the knowledge related to different molecules in pretraining dataset and the knowledge of change in conformation in force field dataset.
If so, the pretraining should help the model to generalize better on unseen molecules.
We use GemNet-T [14] as the backbone model and use the same set of hyperparameters as [14].
The force is computed as the negative gradient of the predicted energy with respect to atom coordinates.
The experiments are conducted with a Nvidia V100 GPU.
4.1.1 In-domain test
In this setting, the training data and test data contains different set of conformational geometries and corresponding force fields for the same molecule.
We train models for each molecule in the
MD17@CCSD dataset [33].
We follow [14] for the train/validation/test split, using 950 samples for training, 50 samples for validation, and 500 samples for test.
As shown in Table 1, the models finetuned from the pretrained model generally perform better than the models trained from scratch.
[Table]
4.1.2 Out-of-domain test
In this setting, we employ the MD17 DFT dataset [13], which contains 10 kinds of molecules.
We construct the training and testing data in a “train-9-test-1" way.
For each kind of molecule to be tested, we use the samples from the other 9 molecules as the training set.
The test set contains 5k randomly chosen samples, and the training set is built by randomly 50k from each training molecule and mixing the combined 450 samples.
For comparison, we define a naïve constant baseline which always predicts zero force.
The reported test error for this baseline is equivalent to the element-wise mean absolute value of the force vectors.
As shown in Table 2, the test error is generally much higher than the in-domain test.
This is expected as the test molecules are not included in the training set.
Benzene and toluene show a relative low loss, probably because benzene ring, their main structural component, has appeared in several molecules (aspirin, azobenzene, salicylic and naphthalene) in the training set.
This implies that the existence of similar substructure in training data significantly help the test performance.
The models finetuned from the pretrained model perform significantly better than the models trained from scratch.
If the pretraining is conducted without force regularization ( α = 1 ), the improvement become less significant.
Force regularization encourage the pretrained model not only to learn the energy for a given geometry, but also the energy gradient at this geometry.
This makes pretraining more relevant to the force field prediction tasks.
[Table]
4.2 Molecular properties
Following previous works [6, 7, 15]
[Table]
we use scaffold splitting a ratio for train/validation/test as 8:1:1.
This splitting method make the molecules in train set do not share molecular scaffold [29] with the molecules in validation or test set.
The tested molecules are therefore not "similar" to the molecules in training dataset
[Table]
and this splitting method is believed to offer a more challenging yet realistic way compared to random splitting.
For each dataset, we run three times to obtain three different scaffold splits.
The average value and standard deviation of the test metrics are reported.
We use EGNN [34] as the backbone model.
If a dataset include more than one tasks (e.g., SIDER contains 27 tasks), we use a multi-task approach, using a single model to predict all tasks simultaneously, similar to [7].
The checkpoint cooresponding to the lowest validation loss is evaluated on the test set.
We observe that pretraining makes training on downstream task reaches the lowest validation loss much faster than training from scratch, reducing the necessary training epochs from about 300 to 30.
As shown in Table [Table] the models finetuned from the pretrained model generally perform much better than the models trained from scratch.
We do not observe negative transfer as Hu et.al.
[5] did with their multi-task supervised pretraining strategies.
Compared to self-supervised pretraining methods (N-GRAM [35], Hu et.al.
[5], GROVER [7] [Table] and MPG [6]), our supervised pretraining method achieve similar or better results [Table]
[Table]
[Table]
4.3 What have the pretrained models learned?
The experiments above show that pretraining on energy prediction tasks improve the model performance on both force fields and molecular properties prediction.
The improvement on the force fields seems not surprising as the energy is closely related to force
[Table]
making the pretraining task relevant to the downstream tasks.
However, the relation between energy and the tested molecular properties is not so obvious.
This makes us wondering what knowledge the pretrained models have learned.
Both the molecular energy and the molecular properties depend on the molecular structure [12, 40].
This motivates us to test whether the pretrained model learns to embed the structural information, which can be then used in downstream tasks for molecular properties prediction.
Following [41], we employ linear probes to analyze representation learned by the pretrained EGNN model.
Each layer in EGNN model output a node representation h ∈ R D for each atom, where D is the hidden dimension.
The node representation can be summed up as the graph representation.
For linear probing, the pretrained model is frozen.
We use a trainable linear layer to map the learned representation to the output space.
We consider the following two categories of diagnostic tasks to test whether the pretrained model learns to embed the molecular structural information.
4.3.1 Reconstructing the input
We firstly test whether the representation produced by the pretrained model can be used to reconstruct the input of EGNN, the atom types and interatomic distance.
Atom charge classification.
We start with a simple test, whether the node representation contains the information about atom charges (e.g. hydrogen, carbon or oxygen).
For a molecule with n atoms, the node representation for the i -th atom, h i , is sent to a linear probe layer to output p i the probability distribution of atom charges.
Interatomic distance prediction.
This task tests whether the pretrained model learns to embed the 3D geometries.
h i is input to a linear probe layer to output ˆ x , the estimated 3D coordinates.
We compare the test error of these three diagnostic tasks with EGNN models of random parameters.
As illustrated in Figure 2, the probing performance of the pretrained model is significantly better than model of random parameters.
This indicates that the pretrained model learns to embed structure information, although it is only trained to predict molecular energy.
For atom type classification task, the test loss using the representation from the first layer is smaller than that of the last layer, as shown in Figure 2.
This is expected as in EGNN, the atom type information is only directly provided to the first layer.
Model of random parameters quickly lose this information as indicated by the significantly higher loss.
In contrast, the pretrain model learned to pass this information to the last layer with smaller loss.
For interatomic distance prediction task, as the atom coordinates information is sent to each EGNN layer, the test loss do not change significantly across different layers.
The results above show that, although only trained to predict molecular energy, the pretrained model learn to embed molecular structural information in the learned representation.
As structural information is important to various downstream tasks, pretraining improve the performance compared to training from scratch.
4.3.2 Identifying substructure
We then test if the pretrained model learns to embed the information necessary to identify the molecular substructures.
Molecular scaffold classification.
We are interested in higher level molecular structure, scaffold.
From the training data, we choose the top 100 frequent scaffold.
With trainable linear layers, parametrized by W Ts h ∈ R D × D and W
Ts p ∈ R 100 × D , The learned representation is used to predict the scaffold class probability distribution p s .
Given labels y scaffold , the probe layers are trained with negative log likelihodd loss.
Molecular motif identification.
Similar to [7], we consider 85 functional groups 2 , such as aliphatic carboxylic acids and H-pyrrole nitrogens, as the molecular motifs.
We test if the pretrained model embed the information to identify whether each of these fragment appear in the molecule.
We formulate this as a multi-task binary classification problem.
With trainable linear layers, parametrized by W Tf h ∈ R D × D , W T negative ∈ R 85 × D , and W T positive ∈ R 85 × D , the learned representation is used to predict the binary classification probability p j for the j -th fragment.
Given labels y motif , the probe layers are trained with negative log likelihodd loss.
Similar to the input reconstruction tasks, the probing performance of the pretrained model is significantly better than model of random parameters, as shown in Figure 2.
For molecular scaffold classification task, using the representation from the last EGNN layer of the pretrained model results in a much smaller test error compared to that of the first layer, as shown in Figure 2.
Scaffold class is not a input of the EGNN model, and have to be predicted based on the molecular global structure.
This implies that the pretrained model learns to gradually obtain the global representation of the molecules at later layers.
Similar trends are observed for molecular fragment identification task.
As the information pass along to deeper layers, the representation embed more global information and is easier for prediction of molecular subtructure tasks.
Conclusion
We proposed and demonstrated a supervised learning pretraining strategy for molecular force fields and properties prediction.
We find that the pretrained model not only learns to perform the energy prediction task, but also embed the molecular structure information.
Experiments show that, compared to training from scratch, fine-tuning the pretrained model can significantly improve the performance.
The demonstration covers seven molecular properties datasets and two molecular force field datasets, including a new zero-shot test.
New state-of-the-art performance are observed for a few tasks.