An out-of-distribution generalization method for molecular representation learning based on graph neural networks

Abstract

The present invention relates to a method for out-of-distribution generalization of molecular representation learning based on graph neural networks, and belongs to the field of deep learning technology. The method uses unsupervised learning and environmental classifier training. The molecular graph is first divided into environments, and then the environmental classification loss is calculated through a GNN encoder and a predictor. Next, node embeddings are extracted and subgraph masks are generated, and the KL divergence loss between the input graph and the subgraph is calculated. The generated subgraph mask is then applied to the GNN encoder to calculate the subgraph prediction loss. Finally, the model parameters are updated through backpropagation, taking into account the environmental classification loss, KL divergence loss, and subgraph prediction loss. The method optimizes the out-of-distribution generalization ability of the model through iterative training, ensuring a balance between subgraph compression, classification tasks, and environmental losses. At the same time, the loss function combines environmental classification, subgraph generation, and classification tasks, thereby improving the performance and generalization ability of molecular representation learning.

Description

Technical Field

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

[0002] Molecular Representation Learning (MRL) is an important bridge between machine learning and chemistry. Its core goal is to encode molecular data into high-dimensional digital vectors using deep learning models, thereby capturing the structural and chemical properties of molecules and supporting downstream tasks such as molecular property prediction and novel drug discovery. MRL has made significant progress in various fields, such as molecular property prediction, the search for novel antibiotics, and drug-drug interaction prediction. However, the dynamic nature of molecular environments and the complexity of data distribution in reality pose challenges to MRL models in terms of out-of-distribution (OOD) generalization, which has become a major bottleneck for their widespread application.

[0003] Current MRL methods are typically based on the independent and identically distributed (iid) assumption, which states that the training data and test data come from the same distribution. However, in practical applications, this assumption is difficult to hold, such as changes in the chemical space in drug development or new data distributions introduced by public health emergencies. This distribution shift problem often leads to decreased model performance, especially in OOD scenarios in unseen data environments. The introduction of Graph Neural Networks (GNNs) provides a new perspective for solving this problem. It can capture key subgraph characteristics through the graph structure of molecules and extract invariant features that remain consistent under different distributions, thereby improving the generalization ability of the model. Invariant learning has also become an important strategy for solving the OOD generalization problem. Its core is to extract feature patterns that have an invariant relationship with labels from different environments.

[0004] Existing MRL methods have two major shortcomings when dealing with distribution shift and OOD issues: First, the lack of diversity in environmental partitioning makes it difficult to effectively eliminate spurious features related to the environment, limiting the model's OOD generalization ability; second, the high dependency on distribution labels, which often require manual annotation, is not only time-consuming and labor-intensive, but also requires extensive domain knowledge and often fails to accurately reflect the underlying distribution characteristics of molecular data. These shortcomings result in poor performance of existing methods in automatically extracting distribution features and improving generalization performance, becoming a major bottleneck in the application of MRL in complex environments. Summary of the Invention

[0005] The purpose of the present invention is to provide a method for out-of-distribution generalization of 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 of environmental partitioning and high dependence on distribution labels.

[0006] To achieve the above-mentioned purpose, the technical solution of the present invention is: a method for out-of-distribution generalization of molecular representation learning based on graph neural networks, first, the environmental reasoning module infers the potential environmental variables of the molecule in an unsupervised manner, uses the environmental classifier in combination with the reconstruction network, and identifies the impact of different environments on the molecular graph by minimizing the lower bound of evidence optimization, providing the model with environmental information to enhance its generalization ability. The goal of the invariant subgraph generation module is to extract the subgraph that is most relevant to the prediction task from the input molecular graph and ensure that the subgraph remains consistent in different environments. By combining environmental reasoning with invariant subgraph generation, the model is able to identify environmental features and extract robust and task-related subgraph representations, thereby maintaining consistent prediction performance on out-of-distribution data.

[0007] The specific steps are:

[0008] Step 1: The environment classifier is combined with the reconstruction network for unsupervised training;

[0009] Step 2: Load the trained environment classifier, divide the molecular graph into environments, input the molecular graph into the GNN encoder and predictor, and calculate the environment classification loss;

[0010] Step 3: 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. The KL divergence loss between the input graph and the subgraph is calculated using the molecules of the subgraph mask to measure the effect of information compression;

[0011] Step 4: Apply the generated subgraph mask to the GNN encoder to obtain the embedding of the subgraph. After inputting the embedding of the subgraph into the predictor, the subgraph prediction loss is calculated.

[0012] Step 5: The environment classification loss, KL divergence loss, and subgraph prediction loss are weighted and summed according to the preset weights to obtain the final loss. The gradient of the loss with respect to the model parameters is then calculated through backpropagation to update the model parameters.

[0013] Step 6: Repeat Step 2 to Step 5 until all training data are traversed and the training of the GNN encoder, subgraph generator, and predictor is completed. Then, the data is predicted and the final classification result is output.

[0014] The Step 1 is specifically as follows:

[0015] Step 1.1: Environment Classifier Aims to predict the potential environment to which the input molecular graph belongs Since there is no environmental label, the environmental classifier needs to rely on the feedback signal provided by the reconstruction network to learn the posterior distribution of the environment by minimizing ELBO. , where ELBO is an optimization objective in variational inference, which is used to minimize the difference between the true posterior distribution and the variational posterior distribution. Given an input molecular graph , through graph neural network Extracting a global representation of the molecular graph , then connect and tags , through the predictor Predicted environmental distribution:

[0016]

[0017] (1)

[0018] in, is the function of normalized probability distribution, Represents the connection graph embedding and its corresponding label .

[0019] Step 1.2: Reconstruct the network Responsible for generating the reconstruction loss term in ELBO and providing necessary information for training the environment classifier. The reconstruction network receives the molecular graph and environmental labels , for each environment Learn an embedding vector , output pair label Conditional probability prediction of , its main function is to provide reconstruction loss rather than directly improve the accuracy of label prediction. , select the corresponding environment embedding from the environment embedding matrix , then, through the graph neural network Extracting a global representation of the molecular graph ,Will Embedded with the environment Connection, through the predictor Predicted labels:

[0020]

[0021]

[0022] (2)

[0023] in, Indicates selecting the corresponding environment embedding from the environment embedding matrix.

[0024] Step 1.3: After using variational approximation, we can derive the model's evidence lower bound loss function:

[0025] (3)

[0026] in, Is the prior distribution of the environment, usually set to uniform distribution or Gaussian distribution. Reconstruction loss Measures the performance of the reconstructed network in a given environment Next, label predictive ability. Divergence term As a regularization term, it encourages approximate distribution Close to the prior distribution Through this combination, the model can more accurately infer the underlying environment, thereby capturing the key features and environmental information in the data. By minimizing ELBO, the parameters of the model can be efficiently trained and the true posterior distribution can be effectively approximated. This approach not only makes the optimization process feasible, but also builds a unified training module that combines the training of the environment classifier with the reconstruction network.

[0027] The Step 2 is specifically as follows:

[0028] Step 2.1: Load the trained environment classifier, which consists of a encoder and an MLP predictor Composition, according to the pre-set number of environments Perform environmental partitioning on a molecular graph:

[0029] (4)

[0030] Where, Represents a specific training data, Indicates the environment label where the data is divided;

[0031] Step 2.2: After completing the molecular graph environment division, input the molecular graph into encoder and predictor Make predictions in:

[0032] (5)

[0033] in, Indicates that the molecular graph is input After the encoder and predictor, the predicted probability is obtained, Represents the pooling operation, which is used to downsample the feature map;

[0034] Step 2.3: Calculate the environment classification loss :

[0035] (6)

[0036] Where, Indicates that the training set Divide into environments All molecular graphs.

[0037] The Step 3 is specifically as follows:

[0038] Step 3.1: For a given molecular graph ,in Represents the adjacency matrix of the graph, describing the connection relationship between atoms, Represents the node feature matrix, which contains the attribute information of each atom. encoder Extract embeddings:

[0039] (7)

[0040] Where, Represents the embedding matrix of all nodes in the molecular graph, which contains the embedding vector of each node ;

[0041] Step 3.2: For specific nodes , invariant subgraph generator First, embed the node Through an MLP computing node The sampling probability of :

[0042] (8)

[0043] Where, 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: (9)

[0044] Step 3.3: After generating the node mask, the edge sampling mask is obtained by averaging the connected nodes. Given two connected nodes and , generate edge mask ,in, and Representation node and The node mask of the entire node mask matrix is recorded as , the edge mask matrix is recorded as ;

[0045] Step 3.3: Use KL divergence to constrain the subgraph size and calculate the KL divergence loss for:

[0046] (10)

[0047] in, Represents the original molecular graph The true distribution of Indicates from The subgraph sampled from ;

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

[0049] The Step 4 is specifically as follows:

[0050] Step 4.1: Apply the edge mask of the subgraph to the GNN encoder To generate new graph embeddings :

[0051] (11)

[0052] Step 4.2: Embed the new graph With node mask Multiply element by element to get the embedding of the subgraph :

[0053] (12)

[0054] Step 4.3: Embed the subgraph into the input predictor Make a prediction:

[0055] (13)

[0056] in, Represents a given subgraph In the case of The predicted probability of

[0057] Step 4.4: Calculate subgraph prediction loss :

[0058] (14)

[0059] in, Indicates the total number of samples.

[0060] The Step 5 is specifically as follows:

[0061] Total loss function Defined as:

[0062] (15)

[0063] in, and is a weight parameter used to adjust the balance between subgraph compression, classification tasks, and environmental losses to ensure the optimization effect of the model in out-of-distribution generalization tasks.

[0064] The beneficial effects of the present invention are: compared with existing molecular characterization technologies, the present invention can accurately infer the potential environment of molecules and extract robust subgraph representations related to the task by combining environmental reasoning and invariant subgraph generation, thereby maintaining high prediction performance when encountering unknown data distributions. BRIEF DESCRIPTION OF THE DRAWINGS

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

[0066] Figure 2 It is a flow chart of the overall framework of the present invention. DETAILED DESCRIPTION

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

[0068] In this example, the EC50 subset of the "DrugOOD" dataset was used. This dataset contains experimental data for a series of drug molecules to evaluate 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 environmental inference model was trained for 25 epochs. Each epoch optimizes the parameters of the environmental classifier by minimizing the lower bound of evidence. In order to ensure the accuracy of environmental classification, the environmental inference module uses a graph neural network as an encoder and combines it with a reconstruction network for optimization. GNN is used to extract global representations from the molecular graph and is trained through different environmental inference modules. The number of layers of GNN is 3 and the hidden layer dimension is 64.

[0069] A molecular representation learning distribution generalization method based on graph neural network. The steps of this method are as follows Figure 1 As shown, the overall framework is as Figure 2 As shown, the specific steps are:

[0070] Step 1: Environment classifier training.

[0071] The environment classifier is combined with the reconstruction network for unsupervised training. The environment classifier extracts global molecular representation using a GNN encoder by minimizing the lower bound of evidence and combines it with labels to perform environment classification.

[0072] Step 2: Environmental division.

[0073] Load the environment classifier trained in Step 1 and use it to classify the input molecular graph into different environments. After classifying the training set into different environments, perform molecular predictions and calculate the environmental classification loss.

[0074] Step 3: Generate invariant subgraph.

[0075] The graph neural network encoder extracts node embeddings from the molecular graph, and the subgraph generation module extracts task-related subgraph representations. While maintaining the graph embedding, the subgraph generator extracts subgraphs based on the sampling probability of the nodes and calculates the KL divergence loss.

[0076] Step 4: Calculate subgraph prediction loss.

[0077] The generated subgraph mask is used to further extract the embedding of the subgraph in the GNN encoder, and the predictor is used to perform task-related classification predictions and calculate the subgraph prediction loss.

[0078] Step 5: Calculate total loss.

[0079] The final training loss is obtained by weighted summing the environment classification loss (Step 2), KL divergence loss (Step 3), and subgraph prediction loss (Step 4). The model parameters are updated through backpropagation.

[0080] Step 6: Complete training and prediction.

[0081] Repeat Step 2 to Step 5 until all training data are traversed and the training of the GNN encoder, subgraph generator, and predictor is completed. Then, the data is predicted and the final classification result is output. For specific experimental results, see Table 1.

[0082] surface Performance on the DrugOOD dataset

[0083]

[0084] Experimental results on the DrugOOD dataset demonstrate that our molecular representation learning method, based on contextual reasoning and invariant subgraph generation, outperforms other molecular representation learning methods designed for out-of-distribution generalization on out-of-distribution data. In particular, our method maintains high prediction accuracy when dealing with changes in drug molecular structure or environmental variations, and demonstrates strong generalization capabilities under unknown data distributions.

[0085] The above describes the specific embodiments of the present invention in detail with reference to the accompanying drawings. However, the present invention is not limited to the above embodiments. Various changes can be made within the knowledge of ordinary technicians in this field without departing from the scope of the present invention.

Claims

1. A method for generalizing molecular representation learning out of distribution based on graph neural networks, characterized by: Step 1: The environment classifier is combined with the reconstruction network for unsupervised training; Step 2: Load the trained environment classifier, divide the molecular graph into environments, input the molecular graph into the GNN encoder and predictor, and calculate the environment classification loss. The molecular graph is the DrugOOD dataset; Step 3: Input the molecular graph into the GNN encoder to extract node embeddings; The embedding is then passed into the subgraph extractor to generate a subgraph mask, and the KL divergence loss between the input graph and the subgraph is calculated using the numerator of the subgraph mask; Step 4: Apply the generated subgraph mask to the GNN encoder to obtain the embedding of the subgraph. After inputting the embedding of the subgraph into the predictor, the subgraph prediction loss is calculated. Step 5: The environment classification loss, KL divergence loss, and subgraph prediction loss are weighted and summed according to the preset weights to obtain the final loss. The gradient of the loss with respect to the model parameters is then calculated through backpropagation to update the model parameters. Step 6: Repeat Step 2 to Step 5 until all training data are traversed and the training of the GNN encoder, subgraph generator, and predictor is completed. Then, the data is predicted and the final classification result is output.

2. The method for generalizing molecular representation learning distribution based on graph neural network according to claim 1, characterized in that: The Step 2 is specifically as follows: Step 2.1: Load the trained environment classifier, which consists of a encoder and an MLP predictor Composition, according to the pre-set number of environments Perform environmental partitioning on a molecular graph: ； （1）； Where, is the embedding of the entire molecular graph, Represents a specific training data, Represents the connection graph embedding and its corresponding label , is the function of normalized probability distribution, Indicates the environment label where the data is divided; Step 2.2: After completing the molecular graph environment division, input the molecular graph into encoder and predictor Make predictions in: （2）; in, Indicates that the molecular graph is input After the encoder and predictor, the predicted probability is obtained, Represents the pooling operation, which is used to downsample the feature map; Step 2.3: Calculate the environment classification loss : （3）; Where, Indicates that the training set Divide into environments All molecular graphs.

3. The method for generalizing molecular representation learning distribution based on graph neural network according to claim 1, characterized in that: The Step 3 is specifically as follows: Step 3.1: For a given molecular graph ,in Represents the adjacency matrix of the graph, describing the connection relationship between atoms, Represents the node feature matrix, which contains the attribute information of each atom. encoder Extract embeddings: （4）; Where, Represents the embedding matrix of all nodes in the molecular graph, which contains the embedding vector of each node ; Step 3.2: For specific nodes , invariant subgraph generator First, embed the node Through an MLP computing node The sampling probability of : （5）; Where, 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, the edge sampling mask is obtained by averaging the connected nodes. Given two connected nodes and , generate edge mask ,in, and Representation node and The node mask of the entire node mask matrix is recorded as , the edge mask matrix is recorded as ; Step 3.3: Use KL divergence to constrain the subgraph size and calculate the KL divergence loss for: (7); in, Represents the original molecular graph The true distribution of Indicates from The subgraph sampled from ; The loss function is based on the KL divergence minimization objective, constraining the subgraph generation distribution and variational distribution The distance is less than the preset value to complete the compression of the input graph.

4. The method for generalizing molecular representation learning distribution based on graph neural network according to claim 1, characterized in that: The Step 4 is specifically as follows: Step 4.1: Apply the edge mask of the subgraph to the GNN encoder To generate new graph embeddings : (8); Step 4.2: Embed the new graph With node mask Multiply element by element to get the embedding of the subgraph : (9) Step 4.3: Embed the subgraph into the input predictor Make a prediction: (10); in, Represents a given subgraph In the case of The predicted probability of Step 4.4: Calculate subgraph prediction loss : (11); in, Indicates the total number of samples.

5. The method for generalizing molecular representation learning distribution based on graph neural network according to claim 1, characterized in that: The Step 5 is specifically as follows: Total loss function Defined as: (12); in, and is a weight parameter used to adjust the balance between subgraph compression, classification task and environmental loss.

Images