A drug molecule generation method based on an equivariant diffusion model

Abstract

The present invention provides a method for generating drug molecules based on an equivariant diffusion model, which relates to the technical field of drug molecule generation and specifically includes the following steps: representing the drug molecule structure as a graph structure; designing an improved rotation, translation, and mirror equivariant autoencoder based on a geometric fully-aware network to encode the atomic coordinates; performing geometric fully latent diffusion based on the diffusion model; adding the loss of the autoencoder and the loss of the diffusion model to obtain the total loss, and optimizing the total loss through gradient descent. The technical solution of the present invention overcomes the problem in the prior art that drug molecules cannot learn the latent representation.

Description

Technical Field

[0001] The present invention relates to the field of drug molecule generation, and particularly to a drug molecule generation method based on an equivariant diffusion model. Background Art

[0002] Drug molecule design and discovery methods mostly rely on experimental screening, theoretical calculations, and the empirical knowledge of chemists. However, these methods often seem inadequate when faced with a vast molecular space and complex biological systems. Traditional molecular discovery methods are essentially time-consuming, and it is impossible for humans to explore the vast chemical space within a limited short time, thus restricting the diversity of molecules that can be explored.

[0003] Deep generative models provide a feasible solution to the above problems. By learning the distribution of various compound molecules through deep generative models, stable and robust molecules can be sampled from the generative model. With the powerful computing power of a computer, this process takes less time and requires no human input, enabling chemical researchers to be liberated from time-consuming and laborious molecular screening work to conduct other more valuable research. There are five popular modeling methods for molecular generative models, namely autoregressive models, variational autoencoders, flow-based models, generative adversarial networks, and diffusion models.

[0004] In recent years, as a new type of generative model, diffusion models have been applied to various generative tasks. Diffusion models define a process of gradually perturbing data with noise and learn to reverse the above process by denoising step by step through a neural network. Existing diffusion models still have some problems. First, the necessary information contained in a molecule includes atomic coordinates, types, atomic charges, and certain physicochemical properties of the molecule, which means that a molecule can be regarded as composed of various discrete or continuous features. Therefore, the diffusion model operating in the atomic feature space is restricted in its learning ability when facing multi-modal features, and it is not optimal to implement a unified Gaussian diffusion framework for multiple modalities. Second, which network parameterization denoising kernel can more accurately predict the noise of the sample is also a problem that needs to be solved by diffusion-based molecular generative models.

[0005] Therefore, there is a need for a drug molecule generation method based on an equivariant diffusion model that can learn the latent representation. Summary of the Invention

[0006] The main object of the present invention is to provide a drug molecule generation method based on an equivariant diffusion model to solve the problem that drug molecules in the prior art cannot learn the latent representation.

[0007] To achieve the above object, the present invention provides a drug molecule generation method based on an equivariant diffusion model, which specifically includes the following steps:

[0008] S1, Represent the drug molecular structure as a graph structure.

[0009] S2, Design an improved rotation, translation, and mirror-equivariant autoencoder based on a geometric fully-aware network to encode atomic coordinates.

[0010] S3, Perform geometric fully latent diffusion based on a diffusion model.

[0011] S4, Add the losses of the autoencoder and the diffusion model to obtain the total loss, and optimize the total loss through gradient descent.

[0012] Furthermore, S1 specifically includes the following steps:

[0013] S1.1, Given a graph , respectively represent the node set and edge set of the graph ; represents the number of nodes in the graph, represents the coordinates of the node set in three-dimensional space.

[0014] S1.2, The node features of and value vector features constitute, Each edge of and value vector features are constituted by, respectively represent the lengths of the scalar features of each node and each edge, then: , where represents the scalar feature of the th node, represents the vector feature of the th node, represents the node and node the scalar feature of the edge between represents the node and node the vector feature of the edge between.

[0015] Furthermore, S2 specifically includes the following steps:

[0016] S2.1, Define geometric fully message passing as:

[0017] ;

[0018] ;

[0019] ;

[0020] ;

[0021] , , ;

[0022] Among them, represents the message between node and node . represents the geometrically complete message passing function, which is the feature of the -th node obtained through layers of geometrically complete perception convolution , is a trainable network, is an aggregation function, represents 's neighbor set, is the message of the aggregated node , , and are intermediate variables, is the geometrically complete frame, is the Euclidean norm.

[0023] S2.2. Update the atomic coordinates in the graph:

[0024] ;

[0025] Among them, is the geometrically complete perception module, is the updated atomic feature, is the -th layer and the -th node's atomic coordinate.

[0026] S2.3. Represent the geometrically complete convolution as:

[0027] ;

[0028] Among them, is the atomic coordinate of the -th layer, is the node scalar feature of the -th layer, is the geometrically complete convolution.

[0029] The geometrically complete convolution conforms to the constraints of rotation, translation, and mirror symmetry, i.e.:

[0030] and ;

[0031] Among them, represents a rotation transformation in Euclidean space, represents a translation transformation.

[0032] Furthermore, step S2 further includes the following steps:

[0033] S2.4, The encoding process of the improved rotation-translation-mirror equivariant autoencoder is:

[0034] ;

[0035] Among them, , are the representations of atomic coordinates and node scalar features in the latent space respectively, , are the reconstructed atomic coordinates and node scalar features, are the atomic coordinates and node scalar features before reconstruction, is noise.

[0036] S2.5, The decoding process of the improved rotation-translation-mirror equivariant autoencoder is:

[0037] .

[0038] S2.6, The loss of the autoencoder is defined as follows:

[0039] .

[0040] Furthermore, S3 specifically includes the following steps:

[0041] S3.1, The diffusion model gradually adds noise to the sample through forward diffusion:

[0042] ;

[0043] ;

[0044] Among them, is a preset hyperparameter and ; is the vector after concatenating atomic coordinates and node scalar features, i.e., the latent feature; is the conditional probability; is the standard normal distribution; is the noise; is the identity matrix.

[0045] S3.2, The diffusion model performs reverse diffusion on the noise variable Denoising is performed to approximate clean samples :

[0046] ;

[0047] Among them, is the mean, is the variance; is the distribution fitted by the diffusion model.

[0048] S3.3. By Bayes' formula:

[0049] ;

[0050] Among them, is the noise randomly sampled from the standard normal distribution, that is .

[0051] S3.4. After the diffusion model denoises the latent features, the decoder restores the features to the geometric space, that is:

[0052] ;

[0053] ;

[0054] Among them, is the output of the noise prediction network, , is a random value sampled from the standard normal distribution, that is , is the decoder of the improved rotation, translation, mirroring equivariant autoencoder.

[0055] S3.5. The loss of the diffusion model is defined as:

[0056] ;

[0057] Among them, is the time.

[0058] Furthermore, S4 specifically includes the following steps:

[0059] S4.1. Add the loss of the autoencoder and the loss of the diffusion model to obtain the total loss:

[0060] .

[0061] S4.2. Optimize through gradient descent.

[0062] The present invention has the following beneficial effects:

[0063] The present invention is a model that satisfies rotation, translation, and mirror-equivariant constraints, and has stronger learning ability for chiral molecules compared to conventional E(3)-equivariant models. The E(3)-equivariant models can hardly distinguish the difference between a molecule and its mirror isomer.

[0064] After being trained on the same dataset as other models, the randomly sampled molecules obtained by the present invention are more robust and effective. Secondly, for the molecules randomly sampled by the present invention, the data distribution of their physical properties shows a high degree of fitting with the attribute distribution of the real data, demonstrating that the present invention has sufficient learning ability for the attribute distribution of real drug molecules. BRIEF DESCRIPTION OF THE DRAWINGS

[0065] In order to more clearly illustrate the specific embodiments of the present invention or the technical solutions in the prior art, the following will briefly introduce the drawings required for the description of the specific embodiments or the prior art. Obviously, the drawings in the following description are some embodiments of the present invention. For those of ordinary skill in the art, without creative efforts, other drawings can also be obtained based on these drawings. In the drawings:

[0066] Figure 1 The flowchart of a method for generating drug molecules based on an equivariant diffusion model of the present invention is shown. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0067] The following will clearly and completely describe the technical solutions of the present invention with reference to the drawings. Obviously, the described embodiments are some, but not all, of the embodiments of the present invention. All other embodiments obtained by those of ordinary skill in the art based on the embodiments of the present invention without creative efforts fall within the scope of protection of the present invention.

[0068] As Figure 1 shown, a method for generating drug molecules based on an equivariant diffusion model is used to solve the problem that the prior art cannot learn the latent representation.

[0069] To achieve the above object, the present invention provides a method for generating drug molecules based on an equivariant diffusion model, which specifically includes the following steps:

[0070] S1, representing the drug molecular structure as a graph structure.

[0071] S2, designing an improved rotation, translation, and mirror-equivariant autoencoder based on a geometric fully-aware network to encode the atomic coordinates.

[0072] S3, performing geometric fully latent diffusion based on the diffusion model.

[0073] S4. Add the loss of the autoencoder and the loss of the diffusion model to obtain the total loss, and optimize the total loss through gradient descent.

[0074] Specifically, S1 specifically includes the following steps:

[0075] S1.1. Given a graph , respectively represent the node set and edge set of the graph ; represents the number of nodes in the graph, represents the coordinates of the node set in three-dimensional space.

[0076] S1.2. The node features of are composed of scalar features and value vector features Each edge of is composed of scalar features and value vector features respectively represent the lengths of the scalar features of each node and each edge, then: , where represents the scalar feature of the th node, represents the vector feature of the th node, represents the scalar feature of the edge between node and node , represents the vector feature of the edge between node and node .

[0077] Specifically, S2 specifically includes the following steps:

[0078] S2.1. Define the geometric complete message passing as:

[0079] ;

[0080] ;

[0081] ;

[0082] ;

[0083] , , ;

[0084] where represents node and the node between the messages, represents the geometrically complete message passing function, which calculates the message of the graph neural network through the vector and scalar features of the nodes and the geometrically complete frame . After layers of geometrically complete perception convolution, the feature of the th node , is a trainable network, is the aggregation function, represents the neighbor set of , and is the message of the aggregated node , and are intermediate variables, is the geometrically complete frame, is the Euclidean norm.

[0085] S2.2. Update the atomic coordinates in the graph:

[0086] ;

[0087] Among them, is the geometrically complete perception module, is the updated atomic feature, is the th layer of the th node's atomic coordinate. By introducing two types of atomic features and the geometrically complete frame as input features, an update that is rotation and translation invariant for atomic scalar features and an update that is rotation equivariant for vector features are provided, and the atomic coordinates are updated based on . After updating the features once through geometrically complete message passing and updating the coordinates once through the geometrically complete perception module, a geometrically complete perception convolution is completed, that is, the geometrically complete convolution layer is composed of geometrically complete message passing and coordinate update modules.

[0088] S2.3. In the above process, the vector feature of the node , the vector feature of the edge and the geometrically complete frame are all calculated from the atomic coordinates of the current layer . In the field of molecular generation, the entire molecular graph is often treated as a fully connected graph, that is, any two different points in a molecular graph are connected. The scalar feature of the edge of the graph can actually be regarded as the fully connected adjacency matrix of the current graph, and all convolutional layers are actually exactly the same. It only provides information for the update of node features. Geometric fully convolutional does not update the features of edges. Therefore, the scalar features of edges are regarded as negligible information in the molecular generation model. In fact, the entire geometric fully convolutional can complete the update of coordinates and features only by relying on the coordinates and scalar features of atoms as inputs. The geometric fully convolutional is expressed as:

[0089] ;

[0090] where is the atomic coordinate of the th layer, is the node scalar feature of the th layer, is the geometric fully convolutional.

[0091] The geometric fully convolutional conforms to the constraints of rotation, translation, and mirror equivariance, that is, translation and rotation equivariance, that is:

[0092] and ;

[0093] where represents the rotation transformation in Euclidean space, represents the translation transformation.

[0094] Compared with the SE(3) (rotation, translation, and mirror equivariant) constraint, the E(3) equivariant constraint (i.e., rotation and translation equivariant constraint) is insensitive to reflection changes due to its equivariance to reflection changes. For molecules that are mirror isomers of each other, the SE(3) equivariant model can accurately capture the differences between the two, while the E(3) equivariant model cannot accurately distinguish between the two.

[0095] Specifically, step S2 further includes the following steps:

[0096] S2.4, the encoding process of the improved rotation, translation, and mirror equivariant autoencoder is:

[0097] ;

[0098] where , are the representations of atomic coordinates and node scalar features in the latent space respectively, , are the reconstructed atomic coordinates and node scalar features, are the atomic coordinates and node scalar features before reconstruction, is noise.

[0099] Actually, after the features are geometrically fully convolutionally encoded, we do not directly output them. Instead, we sample once more with the encoded features themselves as the mean and a variance approaching 0. By introducing a small amount of noise into the latent space, the latent space can be made smoother and more continuous. The decoding process of the improved rotation, translation, and mirror-equivariant autoencoder is as follows:

[0100] .

[0101] S2.6, Loss of the autoencoder Is defined as follows:

[0102] .

[0103] Specifically, thanks to the encoder and decoder of the geometrically complete autoencoder, the present invention can implement a denoising diffusion probability model, i.e., a latent space diffusion model, in a low-dimensional and continuous latent space. The diffusion model models the distribution of the latent representation through a forward diffusion process. In this process, the latent representation is gradually perturbed by Gaussian noise until it satisfies the standard normal distribution. At the same time, the denoising kernel of the diffusion model learns the ability to predict the sample noise. During sampling, the diffusion model samples noise from the standard normal distribution and gradually denoises the sample by virtue of the denoising ability of the denoising kernel to obtain the latent representation of the molecule, and reconstructs the latent representation through the decoding process of the geometrically complete autoencoder. The diffusion model can be described by two Markov chains. The forward Markov chain describes the process of gradually adding noise to the latent representation To obtain The process can be described by the conditional probability The reverse Markov chain, i.e., the denoising process, can be defined by S3 specifically includes the following steps:

[0104] S3.1, The diffusion model gradually adds noise to the sample through forward diffusion:

[0105] ;

[0106] ;

[0107] Among them, Is a preset hyperparameter and , All the parameters of the forward process are preset hyperparameters and are not trainable; Is the vector after splicing the atomic coordinates and the node scalar features, i.e., the latent feature; Is the conditional probability; Is the standard normal distribution; Is the noise; Is the identity matrix.

[0108] S3.2, The diffusion model denoises the noise variable through reverse diffusion to approximate the clean sample :

[0109] ;

[0110] wherein, is the mean, is the variance; is the distribution fitted by the diffusion model. is the output of the neural network, that is, a certain neural network is used to predict the mean of the distribution of forward diffusion. Its variance is default the same as the variance of forward diffusion.

[0111] S3.3, By Bayes' formula:

[0112] ;

[0113] wherein, is the noise randomly sampled from the standard normal distribution, that is . The diffusion model does not have to use a neural network to directly predict 's mean, because all but in its mean are known parameters. Therefore, the diffusion model usually uses a neural network to directly predict the random noise .

[0114] S3.4, The present invention parameterizes the denoising kernel with geometric fully convolution. When using geometric fully convolution as a 3D graph denoiser, the diffusion model not only achieves SE(3) equivariance, but also its geometric completeness enables the denoiser to more accurately predict the noise. During the training of the model of the present invention, a strategy of jointly training the autoencoder and the latent space diffusion is adopted. The drug molecule is first encoded by the encoding process of the autoencoder. After the features are mapped to the latent space by the encoder, the diffusion model learns the distribution of the latent representation and calculates , and at the same time, restores the latent representation to the geometric space to obtain the reconstructed features and calculates the reconstruction loss of the autoencoder. The sum of the two is optimized by gradient descent. For the sampling process, noise is directly sampled from the standard normal distribution in the latent space. The diffusion model denoises the latent features and then the decoder restores the features to the geometric space, that is, the diffusion model denoises the latent features and then the decoder restores the features to the geometric space, that is:

[0115] ;

[0116] ;

[0117] Among them, is the output of the noise prediction network, , is an intermediate variable, is a random value sampled from the standard normal distribution, that is , is the decoder of the improved rotation-translation-mirror equivariant autoencoder.

[0118] S3.5, the loss of the diffusion model is defined as:

[0119] ;

[0120] Among them, is time.

[0121] Specifically, S4 specifically includes the following steps:

[0122] S4.1, adding the loss of the autoencoder and the loss of the diffusion model to obtain the total loss:

[0123] .

[0124] S4.2, optimizing through gradient descent.

[0125] The present invention designs a geometrically complete autoencoder that satisfies the SE(3) equivariant constraint based on geometrically complete perception convolution. With the help of the geometrically complete autoencoder, the present invention realizes latent space diffusion for learning latent representations.

[0126] To verify the method provided by the present invention, a comparative experiment was carried out. As shown in Table 1, the molecular stability, effectiveness, and uniqueness of the method GCLDM proposed by the present invention are relatively high. The experiments shown in Table 2 use the polarizability, dipole moment, and heat capacity as conditional constraints respectively, and let each model generate molecules that meet the conditional constraints. The smaller the numbers in Table 2, the smaller the attributes of the generated molecules and the given attribute constraints, that is, the more in line with the proposed requirements. It can be seen from Table 2 that the average absolute error of the molecular attribute prediction of the present invention is smaller, and the controllability of conditional generation is stronger.

[0127] Table 1 Comparison results of GCLDM and molecular generation baseline methods

[0128]

[0129] Table 2 Comparison results of the average absolute error between the expected value and the predicted value during conditional constraint generation

[0130]

[0131] Certainly, the above description is not a limitation of the present invention, and the present invention is not limited to the above examples. Changes, modifications, additions, or substitutions made by those skilled in the art within the scope of the essence of the present invention should also fall within the protection scope of the present invention.

Claims

1. A method for generating drug molecules based on an equivariant diffusion model, characterized in that, Specifically, it includes the following steps: S1. Represent the drug molecular structure as a graph structure; S2. Based on the geometric fully-aware network, design an improved rotation, translation, and mirror-equivariant autoencoder to encode the atomic coordinates; S3. Based on the diffusion model, perform geometric fully latent diffusion; S4. Add the loss of the autoencoder and the loss of the diffusion model to obtain the total loss, and optimize the total loss through gradient descent; S2 specifically includes the following steps: S2.

1. Define the geometric fully message passing as: ； ； ； ； ， ， ； Among them, represents the message between node and node denotes the geometrically complete message passing function, which is the feature of the th node obtained after layers of geometrically complete perception convolution , is a trainable network, is an aggregation function, represents 's neighbor set, is the message of the aggregated node , , and are intermediate variables, is the geometrically complete frame, is the Euclidean norm; S2.

2. Update the atomic coordinates in the graph: ； Among them, is the geometric complete perception module, is the updated atomic feature, is the th atomic coordinates of the S2.

3. Represent the geometric fully convolution as: ； Among them, is the atomic coordinate of the layer, is the node scalar feature of the layer, is geometric fully convolutional; The geometric fully convolution conforms to the constraints of rotation, translation, and mirror equivariance, that is: and ; Among them, represents a rotation transformation in Euclidean space, represents a translation transformation; Step S2 also includes the following steps: S2.

4. The encoding process of the improved rotation, translation, and mirror-equivariant autoencoder is: ； Among them, , are the representations of atomic coordinates and node scalar features in the latent space respectively, , are the reconstructed atomic coordinates and node scalar features, are the atomic coordinates and node scalar features before reconstruction, is noise; S2.

5. The decoding process of the improved rotation, translation, and mirror-equivariant autoencoder is: ； S2.6, Loss of the autoencoder is defined as follows: 。 2. The method for generating drug molecules based on an equivariant diffusion model according to claim 1, wherein, S1 specifically includes the following steps: S1.1, Given a graph , respectively represent the node set and edge set of the graph ; represents the number of nodes in the graph, represents the coordinates of the node set in three-dimensional space; S1.2 The node features of are composed of scalar features and value vector features. Each edge of is composed of scalar features and value vector features. Let and represent the lengths of the scalar features of each node and each edge, respectively. Then: where represents the scalar feature of the th node, represents the vector feature of the th node, represents the scalar feature of the edge between node and node and represents the vector feature of the edge between node 3. A method for generating drug molecules based on an equivariant diffusion model according to claim 1, wherein S3 specifically includes the following steps: S3.

1. The diffusion model gradually adds noise to the sample through forward diffusion: ； ； Among them, is a preset hyperparameter and , is the vector after splicing the atomic coordinates and the node scalar features, that is, the latent feature; is the conditional probability; is the standard normal distribution; is the noise; is the identity matrix; S3.

2. The diffusion model denoises the noise variable through reverse diffusion to approximate the clean sample : ； Among them, is the mean value, is the variance; is the distribution fitted by the diffusion model; S3.

3. From Bayes' formula: ； wherein, is noise randomly sampled from a standard normal distribution, that is ; S3.

4. After denoising the latent features by the diffusion model, the decoder restores the features to the geometric space, that is: ； ； Among them, is the output of the noise prediction network, , is a random value sampled from the standard normal distribution, that is , is the decoder of the improved rotation, translation, and mirror-equivariant autoencoder; S3.5, Loss of the Diffusion Model It is defined as: ； Among them, is time.

4. A drug molecule generation method based on an equivariant diffusion model according to claim 1, characterized in that S4 specifically includes the following steps: S4.

1. Add the loss of the autoencoder and the loss of the diffusion model to obtain the total loss: ； S4.2, optimize by gradient descent.

Images