Molecular Graph Representation Learning via Heterogeneous Motif Graph Construction

1 INTRODUCTION . Graph neural networks ( GNNs ) have been proved to effectively solve various challenging tasks in graph embedding fields , such as node classification ( Kipf & Welling , 2016 ) , graph classification ( Xu et al. , 2018 ) , and link prediction ( Schlichtkrull et al. , 2018 ) , which have been extensively applied in social networks , molecular properties prediction , natural language processing , and other fields . Compared with hand-crafted features in the molecular properties prediction field , GNNs map a molecular graph into a dimensional Euclidean space using the topological information among the nodes in the graph ( Scarselli et al. , 2008 ) . Most existing GNNs use the basic molecular graphs topology to obtain structural information through neighborhood feature aggregation and pooling methods ( Kipf & Welling , 2016 ; Ying et al. , 2018 ; Gao & Ji , 2019 ) . However , these methods fail to consider connections among molecular graphs , specifically the sharing of motif patterns in the molecular graph . One of the critical differences between molecular graphs and other graph structures such as social network graphs and citation graphs is that motifs , which can be seen as common sub-graphs in molecular graphs have special meanings . For example , an edge in a molecule represents a bond , and a cycle represents a ring . One ground truth that has been widely used in explanation of GNNs is that carbon rings and NO2 groups tend to be mutagenic ( Debnath et al. , 1991 ) . Thus , motifs deserve more attention when designing GNNs for motif-level feature representation learning . To this end , we propose a novel method to learn motif-level feature embedding for molecular graphs . We first extract motifs from molecular graphs and build a motif vocabulary containing all these motifs . Then we construct a heterogeneous motif graph containing all motif nodes and molecular nodes . We can apply GNNs to learn motif-level representations for each molecular graph based on the heterogeneous motif graph . The message passing scheme in a heterogeneous motif graph enables interaction between motifs and molecules , which helps exchange information between molecular graphs . The experimental results show that the learned motif-level embedding can dramatically improve the representation of a molecule , and our model can significantly outperform other stateof-the-art GNN models on a variety of graph classification datasets . 2 HETEROGENEOUS MOTIF GRAPH NEURAL NETWORKS . In this section , we propose a novel method to construct a motif-based heterogeneous graph , which can advance motif-based feature representation learning on molecular graphs . Then , we use two separate graph neural networks to learn atom-level and motif-level graph feature representations , respectively . 2.1 MOTIF VOCABULARY OF MOLECULAR GRAPHS . In molecular graphs , motifs are sub-graphs that appear repeatedly and are statistically significant . Specific to biochemical molecule graphs , motifs can be bonds and rings . Analogously , an edge in the graph represents a bond , and a cycle represents a ring . Thus , we can construct a molecule from sub-graphs or motifs out of a motif vocabulary . To represent a molecule by motifs , we first build a motif vocabulary that contains valid sub-graphs from given molecular graphs . To build the motif vocabulary , we search all molecular graphs and extract important sub-graphs . In this work , we only keep bonds and rings to ensure a manageable vocabulary size . However , the algorithm can be easily extended to include different motif patterns . We then remove all duplicate bonds and rings . Some motifs may appear in most of molecules , which carry little information for molecule representation . To reduce the impact of these common motifs , we employ the term frequency–inverse document frequency ( TF-IDF ) algorithm ( Ramos et al. , 2003 ) . In particular , term frequency measures the frequency of a motif in a molecule , and inverse document frequency refers to the number of molecules containing a motif . We average the TF-IDFs of those molecules that contain a motif as the TF-IDF value of the motif . By sorting the vocabulary by TF-IDF , we keep the most essential motifs as our final vocabulary . Figure 1 illustrates the procedure of building motifs vocabulary . 2.2 HETEROGENEOUS MOTIF GRAPH CONSTRUCTION . Based on the motif vocabulary , we build a heterogeneous graph that contains motif nodes and molecular nodes . In this graph , each motif node represents a motif in the vocabulary , and each molecular node is a molecule . Then , we build two types of edges between these nodes ; those are motifmolecule edges and motif-motif edges . We add motif-molecule edges between a molecule node and motif nodes that represent its motifs . We add a motif-motif edge between two motifs if they share at least one atom in any molecule . In this way , we can build a heterogeneous graph containing all motifs in the vocabulary and all molecules connected by two kinds of edges . Appendix A contains detailed pseudocode of constructing a Heterogeneous Motif Graph . One thing to notice is that different motifs have different impacts . We assign different weights to edges based on their ending nodes . In particular , for edges between a motif node and a molecule node , we use the TF-IDF value of the motif as the weight . For the edges between two motif nodes , we use the co-occurrence information point-wise mutual information ( PMI ) , which is a popular correlation measure in information theory and statistics ( Yao et al. , 2019 ) . Formally , the edge weight Aij between node i and node j is computed as Aij = { PMIij , if i , j are motifs TF-IDFij , if i or j is a motif 0 , Otherwise ( 1 ) The TF-IDF value of an edge between a motif node i and a molecular node j is computed as TF-IDFij = C ( i ) j ( log 1 +M 1 +N ( i ) + 1 ) , ( 2 ) where C ( i ) j is the number of times that the motif i appears in the molecule j , M is the number of molecules , and N ( i ) is the number of molecules containing motif i . The PMI value of an edge between two motif nodes is computed as PMIij = log p ( i , j ) p ( i ) p ( j ) , ( 3 ) where p ( i , j ) is the probability that a molecule contains both motif i and motif j , p ( i ) is the probability that a molecule contains motif i , and p ( j ) is the probability that a molecule contains motif j . We use following formulas to compute these probabilities . p ( i , j ) = N ( i , j ) M , p ( i ) = N ( i ) M , p ( j ) = N ( j ) M , ( 4 ) where N ( i , j ) is the number of molecules that contain both motif i and motif j . Figure 2 provides an example of heterogeneous motif graph construction . Note that we assign zero weight for motif node pairs with negative PMI value . 2.3 HETEROGENEOUS MOTIF GRAPH NEURAL NETWORKS . In this part , we build a HM-GNN to learn both atom-level and motif-level graph feature representations . In Section 2.2 , we construct a heterogeneous motif graph that contains all motif nodes and molecular nodes . Here , we first initiate features for each motif node and molecular node . We use the one-hot encoding to generate features for motif nodes . In particular , each motif node i has a feature vector Xi of length |V | , where V represents the motif vocabulary we obtained in Section 2.1 . Given the unique index i of the motif in the vocabulary , we set Xi [ i ] = 1 and other positions to 0 . For molecular nodes , we use a bag-of-words method to populate their feature vectors . We consider each motif as a word and each molecule as a document . By applying the bag-of-words model , we can obtain feature vectors for molecular nodes . Based on this heterogeneous graph , a heterogeneous graph neural network can be applied to learn the motif-level feature embedding for each molecule in the graph . At the same time , each molecule can be easily converted into a graph by using atoms as nodes and bonds as edges . The original molecule graph topology and node features contain atom-level graph information , which can supplement motif-level information . Thus , we employ another graph neural network to learn the atom-level feature embedding . Finally , we concatenate feature embeddings from two graph neural networks and feed them into a multi-layer perceptron ( MLP ) for prediction . Figure 3 shows an example of our HM-GNN model . 2.4 MULTI-TASK LEARNING VIA HETEROGENEOUS MOTIF GRAPH . This part will show that our heterogeneous motif graph can help with graph deep learning models on small molecular datasets via multi-task learning . It is well known that deep learning methods require a significant amount of data for training . However , most molecular datasets are relatively small , and graph deep learning methods can easily over-fit on these datasets . Multi-task learning ( Caruana , 1997 ) has been shown to effectively reduce the risk of over-fitting and help improve the generalization performances of all tasks ( Zhang & Yang , 2017 ) . It can effectively increase the size of train data and decrease the influence of data-dependent noise , which leads to a more robust model . However , it is hard to directly apply multi-task learning on several molecular datasets due to the lack of explicit connections among different datasets . Based on our heterogeneous motif graph , we can easily connect a set of molecular datasets and form a multi-task learning paradigm . GivenN molecular datasetsD1 , · · · , DN , each datasetDi contains ni molecules . We first construct a motif vocabulary V that contains motifs from N molecular datasets . Here , the motif only needs to be shared in some datasets but not all of them . Then , we build a heterogeneous motif graph that contains motifs in the vocabulary and molecules from all datasets . We employ our HM-GNN to learn both graph-level and motif-level feature representations for each molecule based on this graph . The resulting features of each dataset are fed into a separate MLP for prediction . In this process , the motif nodes can be considered as connectors connecting molecules from different datasets or tasks . Under the multi-task training paradigm , our motif-heterogeneous graph can improve the feature representation learning on all datasets . 2.5 EFFICIENT TRAINING VIA EDGE SAMPLING . In Section 2.2 , we construct a heterogeneous motif graph that contains all motif nodes and molecular nodes . As the number of molecular nodes increases , there can be an issue with computational resources . To address this issue , we propose to use an edge sampler to reduce the size of the heterogeneous motif graph . Due to the special structure of our heterogeneous motif graph that it has two kinds of nodes and two kinds of edges , we can efficiently generate a computational subgraph by using the type of an edge . We show how sampling edges can save computational resources . Most GNNs follow a neighborhood aggregation learning scheme . Formally , the ` -th layer of a GNN can be represented by x ` +1i = f ( x ` i , φ ( { eji , x ` j | j ∈ N ( i ) } ) ) ) , ( 5 ) where x ` +1i is the new feature vector of node i , f is a function that combines the ` -th layer ’ s features and the aggregated features , φ is the function that aggregate all neighbors ’ feature vectors of node i , eji is the weight of edgeji , and N ( i ) is the set of node i ’ s neighbors . This equation shows that the time and space complexity are both O ( |E| ) , where |E| is the number of edges in the graph . This means we can reduce the usage of computational resources by removing some edges from the graph . Thus , we employ an edge sampler that samples edges from the graph . And a sampling rule is that we prioritize motif-molecule edges . To sample edges , we first randomly select some molecular nodes as “ starting ” nodes . We run a breadth-first algorithm to conduct a hop-by-hop exploration on the heterogeneous motif graph starting from these nodes . In each hop , we randomly sample a fixed size of edges based on edge type . Note that the first-hop neighbors of each molecular node are the motif nodes , which play essential roles in our heterogeneous motif graph . Thus , we retain all first-hop edges to ensure effective learning of feature representations for motif nodes . Starting from the second hop , we only sample motif-motif edges to retain as much motif information as possible . Figure 4 shows an example of sampling a sub-graph for a 3-layer HM-GNN . Appendix B contains complete sample rules and pseudocode .

Summary: This paper proposes learning molecule representations by using motif-level information. The authors construct a heterogeneous graph which consists of molecules and motifs, then learn representations of them using graph neural networks. The learned features are concatenated with the molecule feature learned from a traditional atom-level graph neural networks and fed into an MLP for property prediction.