A molecular graph contrastive learning pre-training method based on graph pooling data enhancement

Abstract

The present invention discloses a molecular graph comparative learning pre-training method based on graph pooling data enhancement. The present invention includes the following steps: 1) Pre-training model preparation: designing a dual-tower pre-training model, and selecting two different graph pooling methods as data enhancement means for comparative learning. 2) Executing comparative learning pre-training: using comparative learning to perform molecular graph pre-training. 3) Executing molecular graph downstream tasks: using pre-trained model parameters to initialize the graph classification task model for downstream tasks. Compared with the prior art, the present invention proposes a simple and convenient data enhancement method for graph neural network pre-training for the first time. There is no need to design complex data enhancement pre-processing. Only two different graph pooling methods are needed to realize graph neural network comparative learning pre-training, which can be more conveniently applied to downstream task scenarios of graphs.

Description

Technical Field

[0001] The present invention belongs to the field of graph neural networks and molecular design, and in particular relates to a molecular graph neural network pre-training method based on graph pooling data enhancement. Background Art

[0002] Graph data in certain fields, especially chemical molecular graph datasets, have problems such as a small number of labeled data samples for model training and an unbalanced sample distribution. Models that have not been pre-trained are prone to large errors when trained directly on such datasets. At the same time, with the development of self-supervised pre-training, the pre-training-fine-tuning transfer learning model has enabled models to learn universal representations from five-labeled data through self-supervision without the need for a large amount of high-quality labeled data. Contrastive learning pre-training, as a self-supervised pre-training method, has been applied to molecular graph pre-training. However, data augmentation methods used for contrastive learning pre-training usually require more complex data pre-processing, some of which require data pre-processing during batch loading. In addition, some methods will change the structure of the molecular graph, destroying the complete semantics of the molecular graph structure, which is different from the input of downstream tasks. Therefore, designing a method that can be used for graph data augmentation without modifying the graph structure is of great significance and practical value. Summary of the Invention

[0003] The purpose of the present invention is to propose a simpler data enhancement method for molecular graph contrast learning pre-training. This method can perform data enhancement on molecular graph data without modifying or performing other pre-processing on the original pre-training data.

[0004] The technical solution adopted in the present invention is as follows:

[0005] A molecular graph contrastive learning pre-training method based on graph pooling data enhancement, comprising the following steps:

[0006] S1: Obtain a molecular graph dataset and preprocess each molecular graph sample in the dataset to extract its atomic features and chemical bond features as basic features. At the same time, extract the skeleton structure of the molecular graph sample and perform skeleton encoding on the atoms and chemical bonds to form a skeleton feature consisting of atomic skeleton encoding and chemical bond skeleton encoding;

[0007] S2: Construct a pre-trained dual-tower model consisting of two model branches and a contrastive learning projection head. The first branch is composed of a naturally updated online encoder and a first graph pooling layer in cascade, and the second branch is composed of a momentum encoder and a second graph pooling layer using momentum update in cascade. The encoders in both branches use graph neural networks. The input of the pre-trained dual-tower model is basic features and skeleton features, where atomic features and atomic skeleton encodings are each added after passing through an embedding layer to form an atomic embedding, and chemical bond features and chemical bond skeleton encodings are each added after passing through an embedding layer to form a chemical bond embedding. Atomic embedding and chemical bond embedding are simultaneously used as inputs to the encoders in both branches, while the first and second graph pooling layers in the two branches use different graph pooling algorithms. Therefore, the same input obtains different molecular graph representations with different information content in the two branches. Finally, the contrastive learning loss is generated by inputting different molecular graph representations into the contrastive learning projection head.

[0008] S3: Using the molecular graph dataset preprocessed in S1, pre-train the pre-trained dual-tower model using contrastive learning. By calculating the contrastive learning loss, the model learns the universal distribution representation of molecular graphs, and saves the final pre-trained online encoder and the first graph pooling layer.

[0009] S4: Connect the molecular property prediction head after the pre-trained online encoder and the first graph pooling layer in S3 to build a molecular property prediction model; after fine-tuning the molecular property prediction model on the molecular property prediction task dataset, use it for the molecular property prediction task.

[0010] Preferably, in S1, the atomic characteristics include atomic type, chirality, degree, charge, number of hydrogen atoms, hybridization type, number of valence electrons, whether it is aromatic and whether it is on a ring; the chemical bond characteristics include chemical bond type, whether it is conjugated and stereoisomer type.

[0011] Preferably, the skeleton structure is extracted from the molecular graph using rdkit.

[0012] Preferably, the skeleton codes of the atoms and chemical bonds are represented by numbers, wherein the atomic skeleton code is represented by two labels, label 0 represents the atom on the side chain, and label 1 represents the atom on the molecular skeleton; the chemical bond skeleton code is represented by three labels, label 0 represents the chemical bond on the side chain, label 1 represents the chemical bond on the skeleton, and label 2 represents the chemical bond connecting the skeleton and the side chain.

[0013] Preferably, in the pre-trained dual-tower model, the online encoder and the momentum encoder both adopt the GIN model, and the first graph pooling layer connected to the online encoder adopts average pooling as the graph pooling algorithm, and the second graph pooling layer connected to the momentum encoder network adopts maximum pooling as the graph pooling algorithm.

[0014] Preferably, when performing contrastive learning pre-training in S3, the network parameters of the online encoder are updated by back-propagation gradients in each training, while the momentum encoder does not participate in back-propagation gradient update during the training process. It is only updated by the momentum encoder network parameters of the previous batch and the momentum of the online encoder network parameters before each batch of input. The parameter update method is:

[0015] θ m =mθ m +(1)θ o

[0016] Among them, θ m is the momentum encoder network parameter, θ o are the parameters of the online encoder network, and m is the momentum coefficient.

[0017] As an advantage, when performing contrastive learning pre-training on the pre-trained dual-tower model in S3, the contrastive learning loss is l c The calculation formula is as follows:

[0018]

[0019] Where: f represents the contrastive learning projection head, sim (is the cosine similarity calculation function, and g is the contrastive learning loss temperature hyperparameter; oi and g mi is the output i-th positive sample pair, where g oi is the graph representation output by the online encoder and the graph pooling operation combined with it, g mi is the graph representation output by the momentum encoder model and the graph pooling operation combined with it; g oi and g mj is a set of negative sample pairs output; n is the number of sample pairs contained in each batch.

[0020] Preferably, the molecular property prediction head adopts an MLP network.

[0021] Compared with the prior art, the present invention has the following beneficial effects:

[0022] Compared with the existing technology, the molecular graph pre-training method based on graph pooling data enhancement adopted in the present invention uses different graph pooling methods to perform data enhancement on the final graph representation. On the one hand, this method reduces the data enhancement preprocessing in the pre-training process and speeds up the training speed. On the other hand, it eliminates the difference between the pre-training input and the downstream task input, making the model migration more natural. BRIEF DESCRIPTION OF THE DRAWINGS

[0023] Figure 1 Implementation steps for training a pre-trained model for molecular graph contrastive learning based on graph pooling data augmentation.

[0024] Figure 2 This is the effect of molecular skeleton encoding in molecular graph comparative learning pre-training based on graph pooling data enhancement.

[0025] Figure 3 Pre-training model architectures for molecular graph comparative learning with graph pooling-based data augmentation. DETAILED DESCRIPTION

[0026] The present invention will be further described and illustrated below with reference to the accompanying drawings and specific embodiments.

[0027] like Figure 1 As shown, in a preferred embodiment of the present invention, a molecular graph contrast learning pre-training method based on graph pooling data enhancement is provided, and its steps are shown as S1 to S4:

[0028] S1: Obtain a molecular graph dataset and preprocess each molecular graph sample in the dataset to extract its atomic features and chemical bond features as basic features. At the same time, extract the skeleton structure of the molecular graph sample and perform skeleton encoding on the atoms and chemical bonds to form a skeleton feature consisting of atomic skeleton encoding and chemical bond skeleton encoding.

[0029] The molecular graph dataset in this embodiment is a large-scale unlabeled dataset. The data needs to be screened in advance to delete samples with unreasonable structures. Then the SMILES dataset of the molecule needs to be converted into an undirected graph representation, and the molecular skeleton coding needs to be added. The data preprocessing uses the atomic features and chemical bond features provided by the ogb benchmark dataset. The specific atomic features can include 9 types, including atomic type, chirality, degree, charge, number of hydrogen atoms, hybridization type, number of valence electrons, whether it is aromatic and whether it is on the ring. The chemical bond features can include 3 types, including chemical bond type, whether it is conjugated and stereoisomer type. The skeleton coding of atoms and chemical bonds requires first extracting the skeleton structure of the molecular graph, and skeleton coding the atoms and chemical bonds according to the skeleton structure. The skeleton coding of atoms and chemical bonds is represented by the numbers 0, 1, and 2. As shown in the figure, Figure 2As shown, the molecular skeleton can be extracted according to rdkit, and the atomic skeleton code and the chemical bond skeleton code are encoded with labels. The atomic skeleton code is represented by two labels, label 0 represents the atom on the side chain, and label 1 represents the atom on the molecular skeleton; the chemical bond skeleton code is represented by three labels, label 0 represents the chemical bond on the side chain, label 1 represents the chemical bond on the skeleton, and label 2 represents the chemical bond connecting the skeleton and the side chain. The molecular skeleton can be used as a priori knowledge to classify molecules, and the introduction of the molecular skeleton can help the model learn general molecular representations. The present invention does not require data enhancement at the data preprocessing level, which can reduce the memory time during pre-training, and can retain all the features of the molecular graph during the training process for the molecular graph, which can reduce the information loss of the molecular graph semantics during subsequent comparative learning pre-training.

[0030] S2: Construct a pre-trained dual-tower model consisting of two model branches and a contrastive learning projection head. The first branch consists of a naturally updated online encoder cascaded with the first graph pooling layer, and the second branch consists of a momentum encoder cascaded with the second graph pooling layer using momentum updates. The encoders in both branches use graph neural networks. The input to the pre-trained dual-tower model is basic features and skeleton features. Atom features and atomic skeleton encodings are each passed through an embedding layer and added together to form an atom embedding. Chemical bond features and chemical bond skeleton encodings are each passed through an embedding layer and added together to form a chemical bond embedding. Atom embeddings and chemical bond embeddings serve as input to the encoders in both branches. The first and second graph pooling layers in the two branches use different graph pooling algorithms. This allows the same input to produce different molecular graph representations with different information content in the two branches. Finally, the contrastive learning loss is generated by inputting these different molecular graph representations into the contrastive learning projection head.

[0031] In this embodiment, after the two molecular features of basic features and skeleton features are input into the pre-trained double-tower model, four embedding layers are required to obtain the model input. On the basis of the original atomic feature embedding layer AtomEmbedding and the chemical bond feature embedding layer BondEmbedding, the atomic skeleton encoding embedding layer AtomScaffoldEmbedding and the chemical bond skeleton encoding embedding layer BondScaffoldEmbedding are added. The final encoder model input embedding representation is as follows:

[0032] atom=AtomEmbedding(x)

[0033] bond=BondEmbedding(edge)

[0034] atom sc =AtomScaffoldEmbedding(x sc )

[0035] bond sc =BondScaffoldEmbedding(edge sc )

[0036] Therefore, the atomic features, chemical bond features, atomic skeleton codes and chemical bond skeleton codes are embedded in the embedding layer to obtain four feature embeddings: atom, bond, atom, and sc 、bond sc The final atom embedding as the encoder input is represented by the sum of the atom embedding and the atom skeleton embedding atom′=atom+atom sc , chemical bond embedding is represented by the sum of chemical bond embedding and chemical bond skeleton embedding bond′=bond+bond sc In the formula, x represents the input atomic features, edge represents the input chemical bond features, and x sc Represents atomic skeleton encoding, edge sc Represents the chemical bond skeleton code.

[0037] The network structure of the dual-tower model used for pre-training in the present invention is as follows: Figure 3 As shown in the figure, it includes an online encoder network and a momentum encoder network. Two different graph pooling methods are used for the online encoder and momentum encoder as data augmentation methods. Different pooling algorithms ultimately obtain different graph representations, that is, the same molecule has two different representations. For example, the graph representation obtained by maximum pooling is g = Max(h u ),u∈N(u), the molecular graph obtained by average pooling is represented by g=Mean(h u ),u∈N(u), where h u Refers to the molecular graph atomic node representation encoded by the encoder network, and g is the final molecular graph representation. There is a difference in the amount of information in these two graph representations themselves. The principle is the same as that of data enhancement. Without changing the semantics of the input sample, there is no need to perform pre-processing data enhancement at the data level. Pre-training by contrastive learning can bring the molecular representations of the same structure in the sample space closer and the representations of different molecules farther apart. Therefore, based on the momentum encoder network and the online encoder network connected to different graph pooling algorithms in the present invention, although the data samples contained in each input batch are the same and the atomic node representations represented by the model are the same, the molecular graph representations g obtained by different pooling operations are different, which can be naturally used for molecular graph data enhancement.

[0038] In this embodiment, both the online encoder network and the momentum encoder network use the GIN model, where the first graph pooling layer connected to the online encoder uses average pooling as the graph pooling algorithm, and the second graph pooling layer connected to the momentum encoder network uses maximum pooling as the graph pooling algorithm. In this embodiment, the calculation process of the above-mentioned dual-tower model is as follows:

[0039] h o =GIN o (atom′,bond′)

[0040] h m =GIN m (atom′,bond′)

[0041] g o =Mean(h o ),u∈N(u)

[0042] g m =Max(h m ),u∈N(u)

[0043] where h o 、g o The molecular graph node representation calculated by the online encoder network and the graph representation calculated using the node representation after average pooling; h m 、g m The node representations of the molecular graph computed by the momentum encoder network and the graph representation computed using the node representations through max pooling. Because max pooling and average pooling select different features for the node representations and retain different amounts of information, they can be used as a data augmentation method for contrastive learning. Before training, the contrastive learning pre-trained dual-tower model must be initialized.

[0044] However, it should be noted that the two image pooling methods used in the first image pooling layer and the second image pooling layer in this embodiment are not fixed settings of the present invention, but are only the scheme adopted in this embodiment. In fact, other image pooling algorithms can also be used, as long as the image pooling algorithms of the two branches are kept different.

[0045] S3: Using the molecular graph dataset preprocessed in S1, the pre-trained dual-tower model is pre-trained using contrastive learning. By calculating the contrastive learning loss, the model learns the universal distribution representation of molecular graphs, and saves the final pre-trained online encoder and the first graph pooling layer.

[0046] In this embodiment, contrastive learning pre-training is performed, and each batch updates the online encoder network parameters θ by gradient backpropagation o , update the parameters θ of the momentum encoder network by momentum update m, and finally save the pre-trained online encoder network parameters θ o Used for downstream tasks such as molecular graph prediction.

[0047] Specifically, the network parameter update methods of the momentum encoder and the online encoder are different during comparative learning pre-training. The network parameters of the online encoder are updated by back-propagation gradients in each training session, while the momentum encoder does not participate in back-propagation gradient updates during the training process. It is only updated by the momentum encoder network parameters of the previous batch and the momentum of the online encoder network parameters before each batch of input. In other words, in this embodiment, the momentum encoder network model parameters are updated by the momentum of the online network parameters, and are not back-propagated with the online encoder during training. The specific parameter update method used by the momentum encoder is:

[0048] θ m =mθ m +(1)θ o

[0049] Among them, θ m is the momentum encoder network parameter, θ o is the online encoder network parameter, and m is the momentum coefficient. Momentum updates ensure that the molecular graph node representations are similar, and data augmentation is performed by using graph pooling.

[0050] In this embodiment, the online encoder network updates the model parameters by backpropagation, and can learn a more general distribution in the molecular representation space after contrastive learning pre-training based on graph pooling data enhancement. The loss function used for contrastive learning pre-training is as follows:

[0051]

[0052] Where sim is the cosine similarity calculation function, function f is the contrastive learning projection head pre-trained for contrastive learning (in this embodiment, the contrastive learning projection head can use a two-layer MLP network with a relu activation function in the middle), τ is the hyperparameter temperature coefficient of the contrastive learning loss function, and g is the contrastive learning loss function. oi and g mi is the output i-th positive sample pair, where g oi is the graph representation output by the online encoder and the graph pooling operation combined with it, g mi is the graph representation output by the momentum encoder model and the graph pooling operation combined with it; g oi and g mj is the output set of negative sample pairs; n is the number of sample pairs contained in each batch. Optimizing the contrastive learning loss function can shorten the distance between the representations of positive samples in the representation space and increase the distance between the representations of negative samples in the representation space.

[0053] In this embodiment, in each round of pre-training, the online encoder is saved through contrastive learning pre-training for downstream molecular property prediction tasks.

[0054] S4: Connect the molecular property prediction head after the pre-trained online encoder and the first graph pooling layer in S3 to build a molecular property prediction model; after fine-tuning the molecular property prediction model on the molecular property prediction task dataset, use it for the molecular property prediction task.

[0055] In this embodiment, the encoder model used for the downstream molecular property prediction task is the same as the online encoder network model described in S2, and the same hyperparameter initialization model is used. In order to ensure the consistency between the pre-training task and the downstream task, the graph pooling operation connected after the GIN encoder in the molecular property prediction model of the downstream task in this embodiment also uses average pooling as the graph pooling method.

[0056] In this embodiment, the molecular property prediction model uses the online encoder model parameters stored in S3 to initialize the encoder of the molecular property prediction model, and is fine-tuned on the downstream molecular property prediction task in combination with the molecular property prediction head, wherein the molecular property prediction head output is determined by the requirements of the specific task based on the differences in molecular property prediction tasks. In this embodiment, the molecular property prediction head can use an MLP network. Because the molecular property prediction model is initialized using the encoder model parameters learned in the universal representation space, the model can converge to the optimal state more quickly and achieve better results.

[0057] It can be seen that in the above steps S1 to S4, the molecular graph comparative learning pre-training method based on graph pooling data enhancement introduced in the present invention can enable the pre-training model to learn more general representations, while simplifying the data enhancement method of molecular graph comparative learning pre-training, making comparative learning pre-training more general.

[0058] The molecular graph contrast learning pre-training method based on graph pooling data enhancement described in the above embodiments S1 to S4 is applied to a specific example to demonstrate its specific implementation method and technical effect.

[0059] Example

[0060] This embodiment is mainly applied to a set of seven molecular property prediction task datasets included in the benchmark dataset. Each task contains a set of molecular property prediction tasks consisting of several binary classification tasks, including molecular properties used for pharmaceuticals such as toxicity. Details are shown in Table 1.

[0061] Table 1. Basic information of molecular property prediction benchmark dataset

[0062]

[0063] like Figure 1 As shown, in this embodiment, each task is performed according to the method flow described in S1 to S4 (denoted as ours +CL ) and fine-tuned and tested on downstream task data, and the final model prediction results were evaluated by the ROCAUC indicator. In Table 2, this embodiment is fine-tuned and evaluated in seven data sets included in the benchmark data set. Among them, GIN is a graph isomorphism network that has not been pre-trained, PretrianGNN is a GIN trained on a self-supervised semantic prediction pre-training task, GraphMAE is a GIN pre-trained by masked autoencoder generative method, GraphCL is a GIN pre-trained by contrastive learning with pre-processed data enhancement, D-MPNN is the result of a group of neural networks after ensemble learning, grover is a graph transformer network pre-trained by self-supervised semantic prediction, and gps is a model combining a graph neural network and a transformer network. Compared with these baseline models, the method of the present invention has achieved better performance in most data sets (BBBP, Tox21, ToxCast, HIV, BACE) and is in a leading position in all data sets. More importantly, the method of the present invention adopts a simpler data enhancement method than these models based on self-supervised pre-training.

[0064] Table 2. Results of fine-tuning the pre-trained model on molecular property prediction tasks

[0065]

[0066] The above-described embodiment is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Persons skilled in the art may make various changes and modifications without departing from the spirit and scope of the present invention. Therefore, any technical solution obtained by equivalent substitution or equivalent transformation falls within the scope of protection of the present invention.

Claims

1. A molecular graph contrastive learning pre-training method based on graph pooling data enhancement, characterized in that: The following steps are involved: S1: Obtain a molecular graph dataset and preprocess each molecular graph sample in the dataset to extract its atomic features and chemical bond features as basic features. At the same time, extract the skeleton structure of the molecular graph sample and perform skeleton encoding on the atoms and chemical bonds to form a skeleton feature consisting of atomic skeleton encoding and chemical bond skeleton encoding; S2: Construct a pre-trained dual-tower model consisting of two model branches and a contrastive learning projection head. The first branch is composed of a naturally updated online encoder and a first graph pooling layer in cascade, and the second branch is composed of a momentum encoder and a second graph pooling layer using momentum update in cascade. The encoders in both branches use graph neural networks. The input of the pre-trained dual-tower model is basic features and skeleton features, where atomic features and atomic skeleton encodings are each added after passing through an embedding layer to form an atomic embedding, and chemical bond features and chemical bond skeleton encodings are each added after passing through an embedding layer to form a chemical bond embedding. Atomic embedding and chemical bond embedding are simultaneously used as inputs to the encoders in both branches, while the first and second graph pooling layers in the two branches use different graph pooling algorithms. Therefore, the same input obtains different molecular graph representations with different information content in the two branches. Finally, the contrastive learning loss is generated by inputting different molecular graph representations into the contrastive learning projection head. S3: Using the molecular graph dataset preprocessed in S1, pre-train the pre-trained dual-tower model using contrastive learning. By calculating the contrastive learning loss, the model learns the universal distribution representation of molecular graphs, and saves the final pre-trained online encoder and the first graph pooling layer. S4: Connect the molecular property prediction head after the pre-trained online encoder and the first graph pooling layer in S3 to build a molecular property prediction model; after fine-tuning the molecular property prediction model on the molecular property prediction task dataset, use it for the molecular property prediction task.

2. The molecular graph contrastive learning pre-training method based on graph pooling data enhancement according to claim 1, characterized in that: In S1, the atomic characteristics include atomic type, chirality, degree, charge, number of hydrogen atoms, hybridization type, number of valence electrons, whether it is aromatic and whether it is on a ring; the chemical bond characteristics include chemical bond type, whether it is conjugated and stereoisomer type.

3. The molecular graph contrastive learning pre-training method based on graph pooling data enhancement according to claim 1, characterized in that: The backbone structure is extracted from the molecular graph using rdkit.

4. The molecular graph contrastive learning pre-training method based on graph pooling data enhancement according to claim 1, characterized in that: The skeleton codes of the atoms and chemical bonds are represented by numbers, wherein the atomic skeleton code is represented by two labels, label 0 represents the atom on the side chain, and label 1 represents the atom on the molecular skeleton; the chemical bond skeleton code is represented by three labels, label 0 represents the chemical bond on the side chain, label 1 represents the chemical bond on the skeleton, and label 2 represents the chemical bond connecting the skeleton and the side chain.

5. The molecular graph contrastive learning pre-training method based on graph pooling data enhancement according to claim 1, characterized in that: In the pre-trained dual-tower model, both the online encoder and the momentum encoder adopt the GIN model, and the first graph pooling layer connected to the online encoder adopts average pooling as the graph pooling algorithm, and the second graph pooling layer connected to the momentum encoder network adopts maximum pooling as the graph pooling algorithm.

6. The molecular graph contrastive learning pre-training method based on graph pooling data enhancement according to claim 1, characterized in that: During the contrastive learning pre-training in S3, the network parameters of the online encoder are updated by back-propagation gradients in each training, while the momentum encoder does not participate in the back-propagation gradient update during the training process. It is only updated by the momentum encoder network parameters of the previous batch and the momentum of the online encoder network parameters before each batch of input. The parameter update method is: i m =mθ m +(1-m)θ o Among them, θ m is the momentum encoder network parameter, θ o are the parameters of the online encoder network, and m is the momentum coefficient.

7. The molecular graph contrastive learning pre-training method based on graph pooling data enhancement according to claim 1, characterized in that: When the pre-trained dual-tower model is subjected to contrastive learning pre-training in S3, the contrastive learning loss l c The calculation formula is as follows: Where: f represents the contrastive learning projection head, sim() is the cosine similarity calculation function, and g is the contrastive learning loss temperature hyperparameter; oi and g mi is the output i-th positive sample pair, where g oi is the graph representation output by the online encoder and the graph pooling operation combined with it, g mi is the graph representation output by the momentum encoder model and the graph pooling operation combined with it; g oi and g mj is a set of negative sample pairs output; n is the number of sample pairs contained in each batch.

8. The molecular graph contrastive learning pre-training method based on graph pooling data enhancement according to claim 1, characterized in that: The molecular property prediction head adopts an MLP network.

Images