Pre-training molecule embedding GNNs using contrastive learning based on scaffolding
Systems and methods are provided for generating a training dataset for training a molecule embedding module using contrastive learning, wherein the definition of similarity is based on molecular scaffold similarity. For example, systems access a molecular dataset and separate the molecular dataset into positive samples and negative samples. Systems then generate a training dataset comprising the positive samples and negative samples. Systems and methods are also provided for using the trained molecule embedding module to generate molecule embeddings and for building an end-to-end machine learning model configured to perform molecular embedding analysis and molecular property prediction, the model comprising the trained molecule embedding module and a property prediction module.
1. A method for generating training batches for training a molecule embedding module configured to be used in an end-to-end model configured to predict molecular properties, the method executed by one or more computer processors comprising:
accessing a molecular dataset comprising a plurality of molecules; and
separating the molecular dataset into a plurality of positive samples and a plurality of negative samples by at least:
accessing a molecular scaffold dataset comprising a plurality of molecular scaffolds, each molecular scaffold defining a core structure of a particular molecule and one or more underlying characteristics of the particular molecule;
determining that at least two molecules included in the plurality of molecules are similar based on the at least two molecules corresponding to a same molecular scaffold;
labeling the at least two molecules as a positive sample;
determining that at least two other molecules included in the plurality of molecules are dissimilar based on the at least two other molecules corresponding to different molecular scaffolds;
labeling the at least two other molecules as a negative sample; and
generating a training batch comprising the plurality of positive samples corresponding to pairs of similar molecules and the plurality of negative samples corresponding to pairs of dissimilar molecules, the training batch configured to train the molecule embedding module using contrastive learning;
applying the molecule embedding module to a training dataset to train the molecule embedding module to generate molecule embeddings using the contrastive learning;
applying the molecule embedding module to an input molecular dataset comprising a plurality of molecules; and
generating a plurality of molecule embeddings, each molecule embedding corresponding to a molecule included in the input molecular dataset.
2. The method of claim 1 , further comprising:
identifying a target molecule to be added to the training batch;
identifying a particular molecular scaffold associated with the target molecule;
determining that no other molecule included in the training batch is associated with the particular molecular scaffold; and
in response to determining that no other molecule included in the training batch is associated with the particular molecular scaffold, adding the target molecule to the training batch.
3. The method of claim 1 , further comprising:
identifying a target molecule to be added to the training batch;
identifying a particular molecular scaffold associated with the target molecule;
determining that at least one other molecule included in the training batch is associated with the particular molecular scaffold; and
in response to determining that at least one other molecule included in the training batch is associated with the particular molecular scaffold, refraining from adding the target molecule to the training batch.
4. The method of claim 1 , wherein the molecule embedding module comprises one or more graph neural network layers.
5. A method for identifying positive samples from a set of molecular in order to map similar molecules close to each other in a molecular embedding space, the method executed by one or more computer processors comprising:
(i) accessing a molecule embedding module that has been previously trained to map similar molecules close to each other, and dissimilar molecules far apart, the molecule embedding module having been previously trained by a process that includes at least:
accessing a molecular dataset comprising a plurality of molecules; and
separating the molecular dataset into a plurality of positive samples and a plurality of negative samples by at least:
accessing a molecular scaffold dataset comprising a plurality of molecular scaffolds, each molecular scaffold defining a core structure of a particular molecule and one or more underlying characteristics of the particular molecule;
determining that at least two molecules are similar molecules;
labeling the at least two molecules as a positive sample;
determining that at least two other molecules are dissimilar molecules and labeling the at least two other molecules as a negative sample; and
generating a training batch comprising the plurality of positive samples corresponding to pairs molecules corresponding to a same molecular scaffold and the plurality of negative samples corresponding to pairs of dissimilar molecules which correspond to different molecular scaffolds; and
(ii) applying the molecule embedding module to an input molecular dataset to generate a plurality of molecule embeddings, each molecule embedding corresponding to at least one molecule included in the input molecular dataset, wherein the input molecular dataset comprises a plurality of molecular graphs; and
identifying positive molecular samples and negative molecular samples from the plurality of molecule embeddings, wherein each positive molecular sample is identified based on a similar molecular scaffold and each negative molecular sample is identified based on dissimilar molecular scaffold.
6. A method for building an end-to-end machine learning model configured to predict molecular properties, the method executed by one or more computer processors comprising:
(i) accessing a molecule embedding module that has been previously trained to map similar molecules close to each other and dissimilar molecules far apart in an embedding space, by applying the molecule embedding module to a training dataset, the training dataset being generated by a process that includes at least:
accessing a molecular dataset comprising a plurality of molecules;
separating the molecular dataset into a plurality of positive samples and a plurality of negative samples by at least:
accessing a molecular scaffold dataset comprising a plurality of molecular scaffolds, each molecular scaffold defining a core structure of a particular molecule and one or more underlying characteristics of the particular molecule,
determining that at least two molecules are similar molecules,
labeling the at least two molecules as a positive sample,
determining that at least two other molecules are dissimilar molecules and labeling the at least two other molecules as a negative sample; and
generating a training batch comprising the plurality of positive samples corresponding to pairs of molecules corresponding to a same molecular scaffold and the plurality of negative samples corresponding to pairs of dissimilar molecules which correspond to different molecular scaffolds; and
(ii) accessing a property prediction module configured to predict one or more molecular properties; and
(iii) generating an end-to-end machine learning model configured to predict molecular properties by compiling the molecular embedding module with the property prediction module;
applying the molecule embedding module to an input molecular dataset comprising a plurality of molecules;
generating a plurality of molecule embeddings, each molecule embedding corresponding to a molecule included in the input molecular dataset;
identifying a particular molecular property associated with one or more candidate molecules; and
using the property prediction module to predict that a target molecule is associated with the particular molecular property based on a similarity score to the one or more candidate molecules.
7. The method of claim 6 , wherein the molecule embedding dataset comprises a plurality of molecular graphs.
8. The method of claim 6 , wherein the molecule embedding graph comprises one or more graph neural networks.