Molecular property prediction method based on attention graph element pooling

Abstract

This invention discloses a molecular property prediction method based on attention-based primitive pooling. The method includes: acquiring the target molecular structure and constructing a molecular graph, where nodes represent atoms and edges represent chemical bonds; predefining chemically meaningful primitive types, including three-membered rings, four-membered rings, five-membered rings, six-membered rings, and 3-star structures; encoding atomic features using a graph neural network to integrate neighboring chemical information into atomic nodes; enumerating primitive instances isomorphic to the primitive type in the molecular graph; generating an initial representation vector for each primitive instance through average pooling; calculating an assignment matrix from atoms to functional group supernodes using a cross-attention mechanism, with atomic features as queries and primitive instance features as keys; weighted aggregation of atomic features using the assignment matrix to generate supernode features, and reconstructing connections between supernodes based on chemical bonds and shared atomic contributions; progressively pooling the functional group-level molecular graph into a pharmacophore-level molecular graph through hierarchical iteration; and finally outputting the molecular property prediction results through global average pooling and a multilayer perceptron. This invention achieves hierarchical pooling and interpretable representation learning of molecular structures through predefined chemical primitives and cross-attention assignment matrices. It can explicitly identify functional groups and handle cyclic overlapping structures, improving prediction accuracy while providing chemical interpretability. It is suitable for high-throughput screening and candidate molecule evaluation in drug discovery.

Description

Technical Field

[0001] This invention belongs to the field of drug molecule computation and graph deep learning, specifically involving a graph neural network pooling method for molecular property prediction, and particularly a molecular representation learning method based on graph neural network encoding followed by attention mechanism for primitive allocation and hierarchical pooling. Background Technology

[0002] Molecular property prediction is a core task in drug discovery, aiming to predict properties such as solubility, toxicity, and bioactivity from molecular structures, thereby accelerating the screening of candidate drugs. In recent years, Graph Neural Networks (GNNs) have become the mainstream method for molecular property prediction due to their ability to directly process molecular graph structures. In GNNs, graph pooling is a key operation for extracting the overall molecular representation.

[0003] Existing graph pooling methods have the following technical defects: (1) Insufficient functional group recognition ability: Existing pooling methods (such as DiffPool and SAGPool) cluster based on the similarity of atomic features, which cannot effectively identify functional groups with clear chemical significance in molecules (such as benzene rings, carbonyl groups, amino groups, carboxyl groups, etc.). Functional groups are the core units that determine the chemical properties of molecules. Existing methods do not perceive them well enough, resulting in limited prediction accuracy; (2) Difficulty in handling cyclic structures: There are a large number of cyclic structures in molecules (such as aromatic rings and heterocycles). Existing pooling methods treat rings as ordinary node sets, which makes it difficult to capture the overall chemical characteristics of rings. In particular, the problem of assigning shared atoms in fused ring compounds (such as naphthalene and anthracene) cannot be solved by hard clustering; (3) Poor interpretability: Medicinal chemists need to understand which molecular fragments the model is based on for prediction. The clustering results of existing pooling methods lack chemical semantics and cannot provide interpretable prediction basis, which limits their practical application in drug discovery; (4) Limitation of fixed rules: Existing rule-based primitive pooling methods rely on predefined chemical rules (such as distance weight, atom type weight) to calculate and assign weights. These rules are fixed and cannot be learned. They cannot be adaptively adjusted according to specific tasks, which limits the expressive power and generalization ability of the model.

[0004] Therefore, there is an urgent need for a molecular pooling method that can identify functional groups, handle cyclic overlapping structures, has chemical interpretability, and can be learned end-to-end. Summary of the Invention

[0005] This invention addresses the shortcomings of existing technologies by providing a molecular property prediction method based on attention primitive pooling. First, a graph neural network is used to encode atomic features. Then, using chemical primitives in the molecule as basic pooling units, an attention mechanism is employed to achieve a learnable allocation from atoms to primitive instances, constructing a hierarchical molecular representation with chemical semantics, thus enabling interpretable, end-to-end molecular property prediction. The details are as follows: A molecular property prediction method based on attention primitive pooling, comprising the following steps: Step S1, Molecular Graph Construction and Preprocessing: Obtain the structural information of the target molecule and model it as a graph structure. Each node in the graph corresponds to an atom, and each edge corresponds to a chemical bond. Each atom node is assigned an initial feature vector, containing chemical features such as atom type, hybridization state, charge, aromaticity, and number of hydrogen atoms. Step S2, Chemical Primitive Predefinition: Several chemically meaningful primitive types are predefined. Each primitive is defined by a fixed number of nodes and connection method, corresponding to a specific molecular substructure. The primitives include three-membered rings, four-membered rings, five-membered rings, six-membered rings, 3-stars, etc., with the number of nodes limited to 3, 4, 5, or 6. This invention only predefines chemical primitive structures with no more than 5 nodes. Molecular primitives with more than 6 nodes are not included in the protection and implementation scope of this invention. Step S3: Graph Neural Network Encoding of Atomic Features: A graph neural network is used to encode the molecular graph, enabling each atomic node to incorporate the chemical information of its neighboring nodes. Through multi-layer message passing and node updates, atomic features rich in local chemical environments are obtained. Step S4, Molecular Primitive Instance Enumeration: Enumerate all subgraphs isomorphic to the predefined primitive type in the molecular graph; these are called primitive instances. A graph isomorphism detection algorithm (such as the Orca algorithm) is used for enumeration and extraction. Random sampling can be performed to control complexity. Step S5: Generate the representation of primitive instances: For each primitive instance, extract the encoding features of its internal atoms and use average pooling to generate the initial representation vector of the primitive. Step S6: Calculate the attention assignment matrix: Create a functional clique supernode for each primitive instance. Using a cross-attention mechanism, with atomic features as queries and primitive instance features as keys, calculate the weight of each atom belonging to each supernode. A masking mechanism is used to restrict atoms to only calculating attention to their respective primitive instances, and Softmax normalization is applied. Step S7, Functional group supernode feature aggregation: Using the attention allocation matrix, the atomic features are weighted and aggregated to the functional group supernode to generate the feature vector of the supernode; Step S8, Functional Group Supernode Edge Reconstruction: Based on the chemical bond and attention allocation matrix of the original molecular graph, the connection relationships between supernodes are reconstructed. Each chemical bond contributes a weight to the supernode pair, and shared atoms in fused rings are also supplemented through weighting. Step S9, Hierarchical Pooling Iteration: Using the obtained functional group-level molecular graph as input, repeat steps S3 to S8 to generate a higher-level pharmacophore-level molecular graph, thereby realizing the hierarchical construction of molecular representation; Step S10, Molecular Property Prediction: Global pooling is performed on the final molecular graph to obtain the overall molecular representation vector. Average pooling is then applied, and the vector is input into a multilayer perceptron classifier to output the property prediction results. Step S11, Model Optimization and Training: Based on the prediction results calculated in step S10 and the corresponding real data, calculate the loss value of the network model prediction; when the loss value of the network model converges to the minimum or the training reaches the specified number of iterations, the training ends and the trained network model is output as the prediction model. Step S12: Use the prediction model to detect its performance on the test set; This invention utilizes predefined chemical primitives and a cross-attention mechanism to assign atoms to functional group supernodes, achieving hierarchical pooling and representation learning of molecular structures. This method can explicitly identify functional groups and handle cyclic overlapping structures, improving the accuracy of molecular property prediction while providing chemical interpretability. It is suitable for high-throughput screening and candidate molecule evaluation in drug discovery. Attached Figure Description

[0006] Figure 1 This is an overall flowchart of the method of the present invention; Figure 2 This is a visual flowchart of the algorithm of the present invention; Figure 3 This is a schematic diagram of the predefined primitives in step S2 of the present invention, showing the visual definitions of ternary rings, quaternary rings, pentagonal rings, hexagrams, and 3-stars; Figure 4 This is a schematic diagram of the average pooling of primitive instances in step S5 of the present invention; Figure 5-6 This is a schematic diagram of supernode feature aggregation and edge reconstruction in steps S7-S8 of the present invention. Detailed Implementation

[0007] To better understand the technical solution of this invention, the following is in conjunction with the appendix. Figure 1-6 The embodiments of the present invention will be described in detail below. It should be understood that all other embodiments obtained by those skilled in the art without inventive effort are within the scope of protection of the present invention.

[0008] A molecular property prediction method based on attention primitive pooling includes the following steps.

[0009] Step S1, Molecular Map Construction and Preprocessing: Obtain the structural information of the target molecule and construct the molecular map. ,in: For a set of atomic nodes, This represents the total number of atomic nodes, with each node corresponding to one atom. It is a set of chemical bond edges; Represents atoms With atoms There are chemical bonds between them; Atomic feature matrix, i-th row atomic nodes The initial eigenvectors are denoted as Atomic characteristics include chemical features such as atomic type, hybridization state, charge, aromaticity, and number of hydrogen atoms.

[0010] Step S2, Chemical Primitive Predefinition: A primitive refers to a connected subgraph with a specific topological structure, whose number of nodes is a fixed value k (usually k=3, 4, 5). Different node connection methods constitute different primitive types. For example... Figure 3 As shown, the predefined primitive types in this invention include: three-membered ring (triangle), four-membered ring, five-membered ring, six-membered ring, 3-star, etc. These primitive types correspond to substructures with clear chemical significance in molecules, such as cyclopropane, four-membered ring, five-membered ring, benzene ring, etc. Primitive instance: refers to a specific subgraph that actually exists in the input molecular graph and is isomorphic to a certain primitive type; Primitive enumeration: This refers to the process of finding all primitive instances within an input molecular graph. A predefined set of primitive types with chemical meaning is used. ,in This represents the total number of primitive types, including ring-structure primitives and branch-structure primitives. Three-dimensional ring 3 Three atoms bonded to each other Cyclopropane, ethylene oxide Four-element ring 4 Four atoms form a ring β-lactam ring Five-element ring 5 Five atoms form a ring Pyrrole, furan, thiophene Six-membered ring 6 Six atoms forming a ring Benzene ring, pyridine ring, pyrimidine ring 3-star 4 The central atom connects the three leaves tertiary carbon, amino, hydroxyl .

[0011] Step S3: Graph Neural Network Encoding of Atomic Features: A graph neural network is used to encode the molecular graph constructed in step S1, so that each atomic node incorporates the chemical information of its neighboring nodes to obtain the encoded atomic feature matrix. Specifically, let the graph neural network contain... Layer, number layer The general process of message passing and node updates is as follows: in, Atomic nodes in step S1 Initial feature vector; For the first Atomic nodes after layer encoding eigenvectors; The neighborhood feature aggregation function is used, employing summation, averaging, or maximum value operations; The node feature update function is implemented using a graph convolutional network (GCN). go through After layer encoding, the encoded atomic feature matrix is ​​obtained. , of which OK atomic nodes The encoded feature vector is denoted as , The encoded feature dimension; In an embodiment of the present invention, a two-layer graph convolutional network (GCN) is used as the encoder, with each layer followed by a ReLU activation function and a batch normalization layer.

[0012] Step S4, Enumeration of Molecular Primitive Instances: The molecular graph constructed in step S1 Enumerate all sets of predefined primitive types in step S2. Matching connected subgraphs yields a set of primitive instances. ,in: This represents the total number of primitive instances in the first layer. Each primitive instance It is a set of nodes that satisfies ; Graphical Instance Induced subgraphs and a certain primitive type ; The following algorithm is used to enumerate triangular primitives (ternary rings): For enumeration of primitives with 4 or more nodes, the Orca algorithm is used, with a time complexity of O(n log n). ,in This represents the maximum atomic degree in the molecular graph. To control computational complexity, for large-scale molecular datasets, a random sampling strategy can be used, retaining a maximum of [number missing] primitives per type. Each primitive instance.

[0013] Step S5: Generate a representation of primitive instances: For each primitive instance enumerated in step S4 From the atomic feature matrix encoded in step S3 Extract the feature vectors of its internal atomic nodes to obtain the feature vector set. .like Figure 4 As shown, average pooling is used to aggregate the above feature vectors to generate the initial representation of this primitive instance. The aggregation formula is: Then, the representation vectors of all primitive instances are stacked row-wise to obtain the primitive instance feature matrix. , of which OK .

[0014] Step S6: Calculate the attention assignment matrix for each primitive instance. Create a corresponding functional group supernode The total number of supernodes is Using the atomic feature matrix encoded in step S3 and the primitive instance feature matrix obtained in step S5 The assignment matrix is ​​calculated using a cross-attention mechanism. ,in Represents atomic nodes Belonging to the functional group supernode The weights; The core idea of ​​the attention mechanism is to use atomic features as queries and primitive instance features as keys to calculate the degree of attention that atoms give to each primitive. The calculation steps are as follows: 1. Calculate atomic nodes With primitive instances The attention score between them is calculated using the following formula: in: For the atomic nodes in step S2 The encoded feature vector; For the primitive instance in step S5 The representation vector; , For learnable query and key projection matrices; For projection dimensions; 2. Ensure that atoms only calculate attention for their respective primitive instances through a masking mechanism: 3. Obtain the normalized weight assignment using the Softmax function: The assignment matrix satisfies the following properties: (1) Nonnegativity: for all ,have ; (2) Row normalization: for all ,have ; (3) Soft allocation: Allows an atom to be assigned to multiple functional group supernodes with different weights, thus effectively handling the problem of shared atoms in fused ring compounds; In a preferred embodiment of the present invention, a multi-head cross-attention mechanism is used to calculate the allocation matrix: in To focus on the number of heads, in this embodiment , and For a learnable projection matrix, Indicator functions ensure that atoms only compute attention for the primitive instances to which they belong.

[0015] Step S7: Functional group supernode feature aggregation; using the allocation matrix obtained in step S6. Aggregate the atomic features encoded in step S3 to generate a functional group supernode feature matrix: in The OK Indicates a functional group supernode The eigenvectors of are denoted as . The physical meaning of this aggregation operation is that the feature of each functional group supernode is a weighted sum of its internal atomic features, and the weights are determined by the attention allocation matrix in step S6.

[0016] Step S8: Reconstruction of functional group supernode edges. Based on the molecular graph chemical bond edge set constructed in step S1. and the allocation matrix obtained in step S6 Reconstruct the connections between functional group supernodes: 1. First, initialize the supernode adjacency matrix. It is a zero matrix; 2. For example Figure 5 As shown, for each chemical bond Its two endpoint atoms belong to multiple functional group supernodes. This chemical bond pair supernode pair The edge weight contribution is: After summing all chemical bonds, the supernode adjacency matrix is ​​obtained: like Figure 6 As shown, for fused-ring compounds, shared atoms also contribute to the connection between functional groups, as supplemented by the following formula: in Balance coefficient, in the basic embodiment of the present invention ; 3. Symmetrize the adjacency matrix (applicable to undirected molecular graphs): 4. Extract non-zero elements from the symmetric adjacency matrix to obtain the edge set of the functional group-level molecular graph. and : Thus, the functional group-level molecular diagram after pooling was obtained. ,in: A set of functional group supernodes; A set of functional group supernodes; It is the feature matrix of the functional group supernode.

[0017] Step S9: Hierarchical pooling iteration. The functional group-level molecular diagram obtained in step S8 is then processed... As a new input graph, repeat steps S4 to S8 to generate a higher-level pharmacophore-level molecular graph. For the first Layered pooling The output coarsened map is .

[0018] Step S10: Molecular property prediction. The final hierarchical molecular map obtained in step S9 is then used for... Global pooling is performed to obtain the overall representation vector of the molecule. In the basic embodiment of this invention, global average pooling is used, with the following formula: In a preferred embodiment of the invention, attention pooling is used instead of average pooling, allowing the model to automatically focus on functional groups important to the prediction task: in, and These are learnable parameters; Represent the entire molecule as a vector Input the multilayer perceptron classifier (MLP) and output the molecular property prediction results, as shown in the formula: in Let C be the class probability vector, and C be the number of classes.

[0019] Step S11: Model Optimization and Training. Using the labeled molecular property dataset, with the goal of minimizing the loss function between predicted and true values, the model is optimized using backpropagation and gradient descent algorithms, while simultaneously optimizing all learnable parameters from steps S3 to S10. This reduces the loss value for the next prediction, thereby further optimizing the network model. The loss function used is cross-entropy loss, with the formula: The AdamW optimizer is used during training, with an initial learning rate of 0.001 and weight decay of 1e-5. Training ends when the network model's loss value converges to the minimum or the specified number of iterations are reached, and the trained network model is output as the prediction model.

[0020] Step S12: Test the data to be predicted using the trained network model, employing the same prediction method as steps S3-S10. The performance of this model is evaluated using Cross-Entropy Loss, which improves accuracy by approximately 15% compared to existing superior prediction algorithms. The primitive pooling method enhances the interpretability of the deep learning model's prediction algorithm while simultaneously improving prediction accuracy.

[0021] The embodiments described above are merely preferred embodiments of the present invention and are not intended to limit the present invention in any way. Other variations and modifications are possible without departing from the technical solutions described in the claims.

Claims

1. A method for predicting molecular properties based on attention primitive pooling, characterized in that, Includes the following steps: Step S1: Obtain the structural information of the target molecule and construct a molecular map. ,in For a set of atomic nodes, It is a set of chemical bond edges. The atomic characteristic matrix; Step S2: Predefine a set of primitive types with chemical meaning. The primitive types include ring-shaped primitives and branch-shaped primitives; Step S3: Encode the molecular graph using a graph neural network to obtain the encoded atomic feature matrix. ; Step S4: Enumerate all sets of predefined primitive types in the molecular graph. The matching connected subgraphs yield a set of primitive instances. Each primitive instance ; Step S5: For each primitive instance The representation vector of the primitive instance is generated from the encoded features of its internal atomic nodes through a pooling operation. ; Step S6: Create a corresponding functional clique supernode for each primitive instance, and calculate the allocation matrix using a cross-attention mechanism. ,in Represents atomic nodes Belonging to the functional group supernode The weights are calculated using the following formula: S i j = e x p W Q h i e n c ⋅ W K z j T d k ∑ k = 1 m 1 W Q h i e n c ⋅ W K z k T d k ⋅ 1 v i ∈ I j ,in and For a learnable projection matrix, 1 ⋅ As an indicator function, it ensures that atoms only compute attention for the primitive instances to which they belong; Step S7: Aggregate the encoded atomic features through the allocation matrix to generate a functional group supernode feature matrix. ; Step S8: Chemical bond edge set based on the original molecular diagram and the attention allocation matrix The connectivity between functional group supernodes is reconstructed to obtain the functional group-level molecular graph. ; Step S9: Using the functional group-level molecular map as input, repeat steps S3 to S8 at least once to generate a higher-level pharmacophore-level molecular map. ; Step S10: Perform global pooling on the final molecular graph to obtain the overall molecular representation vector, input it into the multilayer perceptron, and output the prediction results of molecular properties. Step S11: Using the labeled molecular property dataset, train the graph neural network, cross-attention parameters, and multilayer perceptron parameters end-to-end with the goal of minimizing the loss function between the predicted and true values.

2. The method according to claim 1, characterized in that, The primitive types mentioned in step S2 include at least one of the following: three-membered ring primitives, four-membered ring primitives, five-membered ring primitives, six-membered ring primitives, and 3-star primitives. This invention only predefines chemical primitive structures with no more than 6 nodes. Molecular primitives with more than 6 nodes are not included in the protection and implementation scope of this invention.

3. The method according to claim 1, characterized in that, The graph neural network mentioned in step S3 is a general graph neural network, such as a graph convolutional neural network (GCN) or a graph attention network (GAT), with at least 2 layers.

4. The method according to claim 1, characterized in that, The pooling operation described in step S5 is average pooling, and the formula is: .

5. The method according to claim 1, characterized in that, The cross-attention mechanism described in step S6 can employ multi-head attention: S i j = 1 M ∑ m = 1 M e x p W Q m h i e n c ⋅ W K m z i T d k ∑ k = 1 m 1 e x p W Q m h i e n c ⋅ W K m z k T d k ⋅ 1 v i ∈ I j Where M is the number of attention heads and .

6. The method according to claim 1, characterized in that, The reconstruction of the supernode connections in step S8 includes: for each chemical bond Its contribution to supernode pairs The edge weight contribution is Regarding the contribution of shared atoms in fused-ring compounds, through... In addition, among them λ ∈ 0 1 For balance coefficient, This indicates the total number of nodes.

7. The method according to claim 1, characterized in that, The hierarchical pooling iteration in step S9 includes: before each pooling operation, performing at least one layer of graph neural network message passing on the coarsened molecular graph of the current layer, so that the supernode features are fused with their neighbor information.

8. The method according to claim 1, characterized in that, The global pooling described in step S10 is average pooling, and the formula is: Alternatively, attention pooling can be used instead of average pooling, allowing the model to automatically focus on functional groups important for the prediction task. , β v = e x p b T W r h v ∑ u ∈ V L e x p b T W r h u ,in, and These are learnable parameters.

9. The method according to claim 1, characterized in that, The molecular properties include at least one of solubility, toxicity, bioactivity, blood-brain barrier penetration, and metabolic stability.

10. A molecular property prediction system based on attention primitive pooling, characterized in that, include: (1) Molecular diagram construction module, used to convert molecular structures into molecular diagrams; (2) The graph neural network GCN encoding module is used to encode atomic features, so that atomic nodes can fuse neighboring chemical information; (3) Primitive Enumeration Module, used to enumerate primitive instances that match the predefined primitive type on the molecular graph; (4) Cross-attention allocation module, used to calculate the allocation weight matrix from atoms to functional group supernodes; (5) Feature aggregation module, used to aggregate atomic features through allocation matrix to generate supernode features; (6) Edge reconstruction module, used to reconstruct the connection relationships between supernodes; (7) Hierarchical pooling module, used for iterative generation of multi-scale molecular representations; (8) Property prediction module, used to output the prediction results of molecular properties.

Images