Molecular representation learning distribution external generalization method based on graph neural network

Abstract

The invention relates to a molecular representation learning distribution external generalization method based on a graph neural network, and belongs to the technical field of deep learning. According to the method, through unsupervised learning and environment classifier training, a molecular graph is firstly subjected to environment division, and then environment classification loss is calculated through a GNN encoder and a predictor. Then, extracting node embedding, generating a sub-graph mask, and calculating KL divergence loss between the input graph and the sub-graph; then, the generated sub-graph mask acts on a GNN encoder, and sub-graph prediction loss is calculated; and finally, comprehensively considering environment classification loss, KL divergence loss and subgraph prediction loss, and updating model parameters through back propagation. According to the method, the distribution external generalization ability of the model is optimized through iterative training, the balance among subgraph compression, classification tasks and environment loss is ensured, meanwhile, the loss function combines environment classification, subgraph generation and classification tasks, and the performance and generalization ability of molecular representation learning are improved.

Description

Technical Field

[0001] The present invention relates to a method for out-of-distribution generalization of molecular representation learning based on graph neural networks, belonging to the technical field of deep learning. Background Art

[0002] Molecular Representation Learning (MRL) is an important bridge connecting machine learning and chemistry. Its core goal is to encode molecular data into high-dimensional digital vectors through deep learning models, so as to capture the structural characteristics and chemical properties of molecules, and provide support for downstream tasks such as molecular property prediction and new drug discovery. MRL has made remarkable progress in many fields, such as molecular property prediction, search for new antibiotics, and prediction of drug-drug interactions. However, the dynamic changes in the molecular environment and the complexity of data distribution in reality pose challenges for out-of-distribution (OOD) generalization of MRL models, which has become an important bottleneck for its wide application.

[0003] Current MRL methods usually assume independent and identically distributed (i.i.d.), that is, the training data and test data come from the same distribution. However, in practical applications, this assumption is difficult to hold, such as the chemical space change in drug development or the new data distribution introduced by sudden public health events. This distribution shift problem often leads to a decline in model performance, especially in the OOD scenario in the unseen data environment. The introduction of Graph Neural Networks (GNNs) provides a new perspective to solve this problem, which can capture key subgraph features through the graph structure of molecules and extract invariant features that remain consistent under different distributions, thus improving the generalization ability of the model. Invariant learning has also become an important strategy for solving the OOD generalization problem, and its core is to extract feature patterns that have an invariant relationship with labels from different environments.

[0004] Existing MRL methods have two main deficiencies in dealing with distribution shift and OOD problems: one is the lack of diversity in environment division, which is difficult to effectively eliminate false features related to the environment and limits the OOD generalization ability of the model; the other is the high dependence on distribution labels, which usually need to be manually annotated, not only time-consuming and laborious, but also require rich domain knowledge, and often cannot accurately reflect the potential distribution characteristics of molecular data. These deficiencies lead to the poor performance of existing methods in automatically extracting distribution features and improving generalization performance, becoming an important bottleneck for the application of MRL in complex environments. Summary of the Invention

[0005] The object of the present invention is to provide an out-of-distribution generalization method for molecular representation learning based on graph neural networks, aiming to solve the OOD problem in molecular representation learning, specifically the technical problems of insufficient diversity in environmental partitioning and high dependence on distribution labels.

[0006] To achieve the above object, the technical solution of the present invention is: an out-of-distribution generalization method for molecular representation learning based on graph neural networks. First, the environmental inference module infers the latent environmental variables of molecules in an unsupervised manner, and uses the environmental classifier combined with the reconstruction network to optimize by minimizing the evidence lower bound, identifying the influence of different environments on the molecular graph, and providing environmental information for the model to enhance its generalization ability. The goal of the invariant subgraph generation module is to extract the subgraph most relevant to the prediction task from the input molecular graph and ensure the consistency of this subgraph under different environments. By combining environmental inference with invariant subgraph generation, the model can identify environmental features and extract robust and task-related subgraph representations, thus maintaining consistent prediction performance on out-of-distribution data.

[0007] The specific steps are as follows:

[0008] Step1: Combine the environmental classifier and the reconstruction network for unsupervised training;

[0009] Step2: Load the trained environmental classifier, partition the molecular graph into environments, then input the molecular graph into the GNN encoder and predictor, and calculate the environmental classification loss;

[0010] Step3: Input the molecular graph into the GNN encoder to extract node embeddings; then pass the embeddings into the subgraph extractor to generate a subgraph mask, and calculate the KL divergence loss between the input graph and the subgraph through the molecules with the subgraph mask, measuring the effect of information compression;

[0011] Step4: Apply the generated subgraph mask to the GNN encoder to obtain the embeddings of the subgraph, and after inputting the embeddings of the subgraph into the predictor, calculate the subgraph prediction loss;

[0012] Step5: Weighted sum the environmental classification loss, KL divergence loss, and subgraph prediction loss according to preset weights to obtain the final loss, and then calculate the gradient of the loss with respect to the model parameters through backpropagation, and further update the model parameters;

[0013] Step6: Repeat Step2 to Step5 until all training data is traversed, complete the training of the GNN encoder, subgraph generator, and predictor, and then predict the data to output the final classification result.

[0014] The specific content of Step1 is as follows:

[0015] Step1.1: Environmental classifier Aim to predict the potential environment to which the input molecular graph belongs , since there are no environment labels, the environment classifier needs to rely on the feedback signal provided by the reconstruction network and learn the posterior distribution of the environment by minimizing the ELBO , where the ELBO is an optimization objective in variational inference used to minimize the difference between the true posterior distribution and the variational posterior distribution. Given the input molecular graph , extract the global representation of the molecular graph through a graph neural network , then connect and the label , and predict the environment distribution through a predictor :

[0016]

[0017]

[0018] where is a function of the normalized probability distribution represents connecting the graph embedding and its corresponding label .

[0019] Step1.2: Reconstruction Network Responsible for generating the reconstruction loss term in the ELBO and providing necessary information for the training of the environment classifier. The reconstruction network receives the molecular graph and the environment label , learns an embedding vector for each environment , and outputs the conditional probability prediction of the label . Its main role is to provide the reconstruction loss rather than directly improving the accuracy of label prediction. According to the environment label , select the corresponding environment embedding from the environment embedding matrix, then, extract the global representation of the molecular graph through a graph neural network , connect with the environment embedding , and predict the label through a predictor :

[0020]

[0021]

[0022]

[0023] where ​​Indicates selecting the corresponding environmental embedding from the environmental embedding matrix.

[0024] Step1.3: After using variational approximation, the evidence lower bound loss function of the model can be derived:

[0025]

[0026] Among them, is the prior distribution of the environment, usually set to a uniform distribution or a Gaussian distribution. The reconstruction loss measures the prediction ability of the reconstruction network for the label under the given environment . The divergence term acts as a regularization term, encouraging the approximate distribution to be close to the prior distribution . Through this combination, the model can more accurately infer the latent environment, thereby capturing the key features and environmental information in the data. By minimizing the ELBO, the parameters of the model can be efficiently trained, and the true posterior distribution can be effectively approximated. This method not only makes the optimization process feasible but also constructs a unified training module that combines the training of the environmental classifier and the reconstruction network.

[0027] The specific content of the said Step2 is as follows:

[0028] Step2.1: Load the trained environmental classifier, which consists of an encoder and an MLP predictor . Divide the molecular graph according to the preset number of environments :

[0029]

[0030] In the formula, represents a specific training data, represents the environmental label into which the data is divided;

[0031] Step2.2: After completing the division of the molecular graph environment, input the molecular graph into the encoder and the predictor for prediction:

[0032]

[0033] Among them, represents the prediction probability obtained after inputting the molecular graph into the encoder and the predictor, Denotes a pooling operation for downsampling the feature map;

[0034] Step2.3: Calculate the environmental classification loss :

[0035]

[0036] In the formula, Denotes all molecular graphs in the training set Partitioned into the environment ;

[0037] The specific content of Step3 is as follows:

[0038] Step3.1: For the given molecular graph , where Denotes the adjacency matrix of the graph, describing the connection relationship between atoms, Denotes the node feature matrix, containing the attribute information of each atom. First, pass through Encoder Extract the embedding:

[0039]

[0040] In the formula, Denotes the embedding matrix of all nodes in the molecular graph, which contains the embedding vector of each node ;

[0041] Step3.2: For the specific node , the invariant subgraph generator First, calculate the sampling probability of node by passing the node embedding through an MLP:

[0042]

[0043] In the formula, Denotes the Sigmoid activation function, mapping the output of the MLP to the interval [0,1]. The Concrete relaxation method is used to sample the node mask from . By introducing the random variable and the temperature parameter , where Uniform(0,1) denotes a random variable uniformly distributed in the interval [0,1], calculate the node mask:

[0044]

[0045] Step3.3: After generating the node mask, obtain the edge sampling mask by averaging the connected nodes. Given two connected nodes and , generate the edge mask , where and represent the node masks of nodes and . The entire node mask matrix is denoted as , and the edge mask matrix is denoted as ;

[0046] Step3.3: Use the KL divergence to constrain the subgraph size and calculate the KL divergence loss as follows:

[0047]

[0048] where represents the true distribution of the original molecular graph , and represents the subgraph sampled from ;

[0049] The loss function is based on the KL divergence minimization objective, constraining the distance between the subgraph generation distribution and the variational distribution to be less than a preset value to complete the compression of the input graph.

[0050] The specific content of the said Step4 is as follows:

[0051] Step4.1: Apply the edge mask of the subgraph to the GNN encoder to generate a new graph embedding :

[0052]

[0053] Step4.2: Multiply the new graph embedding element-wise with the node mask to obtain the embedding of the subgraph :

[0054]

[0055] Step4.3: Input the subgraph embedding into the predictor for prediction:

[0056]

[0057] where represents the output label given the subgraph ​ Prediction probability;

[0058] Step4.4: Calculate the sub - graph prediction loss :

[0059]

[0060] where represents the total number of samples.

[0061] The specific content of Step5 is as follows:

[0062] The total loss function is defined as:

[0063]

[0064] where and are weight parameters, used to adjust the balance between sub - graph compression, classification tasks, and environmental loss, ensuring the optimization effect of the model in out - of - distribution generalization tasks.

[0065] The beneficial effects of the present invention are: Compared with the existing molecular characterization technologies, the present invention can accurately infer the potential environment of molecules and extract robust sub - graph representations related to tasks by combining environmental reasoning and invariant sub - graph generation, so that it can still maintain high prediction performance when encountering unknown data distributions. BRIEF DESCRIPTION OF THE DRAWINGS

[0066] Figure 1 is a schematic diagram of the steps of the present invention;

[0067] Figure 2 is a flowchart of the overall framework of the present invention. DETAILED DESCRIPTION OF THE INVENTION

[0068] The present invention will be further described below in conjunction with the drawings and specific embodiments.

[0069] In this embodiment, the EC50 subset in the "DrugOOD" dataset is used. This dataset contains experimental data of a series of drug molecules for evaluating their inhibitory effects on specific targets (such as proteins). For each molecule, the given features include node features (39 dimensions) and edge features (10 dimensions). According to the experimental requirements, the model sets 8 different potential environments (i.e., k = 8). These environments represent the potential states of drug molecules under different conditions. The environment inference model is trained for 25 epochs. In each epoch, the parameters of the environment classifier are optimized by minimizing the evidence lower bound. To ensure the accuracy of environment classification, the environment inference module uses a graph neural network as an encoder and combines a reconstruction network for optimization. The GNN is used to extract the global representation from the molecular graph and is trained by different environment inference modules. The number of layers of the GNN is 3, and the dimension of the hidden layer is 64.

[0070] An out-of-distribution generalization method for molecular representation learning based on graph neural network, the steps of which are schematically shown as Figure 1 shown, and the overall framework is as Figure 2 shown. The specific steps are as follows:

[0071] Step1: Training of the environment classifier.

[0072] The environment classifier is trained unsupervised in combination with the reconstruction network. The environment classifier extracts the global molecular representation using the GNN encoder by minimizing the evidence lower bound and combines the labels for environment classification.

[0073] Step2: Environment division.

[0074] Load the environment classifier trained in Step1 and use it to divide the input molecular graph into environments. After dividing different environments through the training set, molecular prediction is performed, and the environment classification loss is calculated.

[0075] Step3: Generation of invariant subgraphs.

[0076] Extract the node embeddings of the molecular graph through the graph neural network encoder, and use the subgraph generation module to extract the subgraph representations related to the task. While ensuring the graph embedding, the subgraph generator extracts subgraphs according to the sampling probability of nodes and calculates the KL divergence loss.

[0077] Step4: Calculation of subgraph prediction loss.

[0078] Use the generated subgraph mask to further extract the subgraph embeddings in the GNN encoder, and perform task-related classification predictions through the predictor to calculate the subgraph prediction loss.

[0079] Step5: Calculation of the total loss.

[0080] The environmental classification loss (Step 2), KL divergence loss (Step 3), and subgraph prediction loss (Step 4) are weighted and summed to obtain the final training loss. The parameters of the model are updated through backpropagation.

[0081] Step 6: Complete training and prediction.

[0082] Repeat Steps 2 to 5 until all training data are traversed, complete the training of the GNN encoder, subgraph generator, and predictor, then predict the data and output the final classification result; see Table 1 for specific experimental results.

[0083] Table 1 Performance on the DrugOOD dataset

[0084] The experimental results on the DrugOOD dataset show that the molecular representation learning method based on environmental reasoning and invariant subgraph generation of the present invention performs better on out-of-distribution data than other molecular representation learning methods designed for out-of-distribution generalization. Especially when dealing with changes in drug molecular structures or environments, the present invention can maintain high prediction accuracy and exhibits strong generalization ability under unknown data distributions.

[0085] The specific embodiments of the present invention have been described in detail above in conjunction with the accompanying drawings. However, the present invention is not limited to the above embodiments, and various changes can be made without departing from the spirit of the present invention within the scope of knowledge possessed by those of ordinary skill in the art.

Claims

1. A method for out-of-distribution generalization of molecular representation learning based on graph neural networks, characterized in that: Step 1: Combine an environment classifier and a reconstruction network for unsupervised training; Step 2: Load the trained environment classifier, divide the molecular graph into environments, then input the molecular graph into the GNN encoder and predictor, and calculate the environment classification loss; Step 3: Input the molecular graph into the GNN encoder to extract node embeddings; Subsequently, input the embeddings into a subgraph extractor to generate a subgraph mask, and calculate the KL divergence loss between the input graph and the subgraph for the molecules passing through the subgraph mask; Step 4: Apply the generated subgraph mask to the GNN encoder to obtain the embeddings of the subgraph, input the embeddings of the subgraph into the predictor, and then calculate the subgraph prediction loss; Step 5: Weighted sum the environment classification loss, KL divergence loss, and subgraph prediction loss according to preset weights to obtain the final loss, and then calculate the gradient of the loss with respect to the model parameters through backpropagation, and further update the model parameters; Step 6: Repeat Step 2 to Step 5 until all training data is traversed, complete the training of the GNN encoder, subgraph generator, and predictor, then make predictions on the data, and output the final classification result.

2. The out-of-distribution generalization method for molecular representation learning based on graph neural network according to claim 1, wherein Specifically, Step 2 is as follows: Step 2.1: Load the trained environment classifier, which consists of an encoder and an MLP predictor and partition the molecular graph according to the preset number of environments as follows: ； （1）； In the formula, is the embedding of the entire molecular graph, represents a specific training data, represents connecting the graph embedding and its corresponding label , is the function of the normalized probability distribution, represents the environmental label into which the data is partitioned; Step 2.2: After completing the molecular graph environment division, input the molecular graph into the encoder and the predictor for prediction: （2）； Among them, represents the predicted probability obtained after inputting the molecular graph into the encoder and predictor, represents a pooling operation for downsampling the feature map; Step 2.3: Calculate the environmental classification loss : （3）； In the formula, represents all the molecular graphs obtained by dividing the training set into the environment .

3. The out-of-distribution generalization method for molecular representation learning based on graph neural network according to claim 1, wherein Specifically, Step 3 is as follows: Step 3.1: For the given molecular graph , where represents the adjacency matrix of the graph, describing the connection relationships between atoms, represents the node feature matrix, containing the attribute information of each atom. First, use the encoder to extract the embedding: （4）； In the formula, represents the embedding matrix of all nodes of the molecular graph, which contains the embedding vectors of each node ; Step 3.2: For a specific node , the invariant subgraph generator first embeds the node and calculates the sampling probability of the node through an MLP : （5）； In the formula, Represents the Sigmoid activation function, maps the output of MLP to the interval [0,1], and uses the Concrete relaxation method to Sampling node mask , by introducing the random variable and temperature parameters , where Uniform(0,1) represents a random variable uniformly distributed on the interval [0,1], and the node mask is calculated: （6）； Step 3.3: After generating the node mask, obtain the edge sampling mask by averaging the connected nodes. Given two connected nodes and , generate the edge mask , where and represent the node masks of nodes and , and the entire node mask matrix is denoted as , and the edge mask matrix is denoted as ; Step 3.3: Use the KL divergence to constrain the size of the subgraph and calculate the KL divergence loss which is （7）； Among them, represents the true distribution of the original molecular graph , represents the subgraph sampled from ; The loss function is based on the KL divergence minimization objective to constrain the subgraph generation distribution to be less than a preset value from the variational distribution to complete the compression of the input graph.

4. The out-of-distribution generalization method for molecular representation learning based on graph neural network according to claim 1, wherein Specifically, Step 4 is as follows: Step4.1: Apply the edge mask of the sub - graph to the GNN encoder to generate a new graph embedding : （8）； Step4.2: Embed the new graph element-wise multiply with the node mask to obtain the embedding of the subgraph : （9）； Step4.3: Embed the sub - graph into the input predictor Make a prediction: （10）； Among them, represents the prediction probability of the output label in the case of a given sub-graph ; Step4.4: Calculate the sub-graph prediction loss : （11）； Among them, represents the total number of samples.

5. The out-of-distribution generalization method for molecular representation learning based on graph neural network according to claim 1, wherein Specifically, Step 5 is as follows: Total loss function is defined as: （12）； Among them, and are weight parameters used to adjust the balance among sub-graph compression, classification tasks, and environmental loss.

Images