Motif-based Graph Self-Supervised Learning for Molecular Property Prediction

1 Introduction . The past decade has witnessed the remarkable success of deep learning in natural language processing ( NLP ) [ 5 ] , computer vision ( CV ) [ 13 ] , and graph analysis [ 19 ] . Inspired by these developments , researchers in the chemistry domain have also tried to exploit deep learning methods for moleculebased tasks such as retrosynthesis [ 47 ] and drug discovery [ 8 ] . To preserve the internal structural information , molecules can be naturally modeled as graphs where nodes represent atoms and edges denote the chemical bonds . Recently , some works applied Graph Neural Network ( GNN ) and some of its variants for molecular property prediction and obtained promising results [ 8 , 25 ] . Though GNNs have achieved remarkable accuracy on molecular property prediction , they are usually data-hungry , i.e . a large amount of labeled data ( i.e. , molecules with known property data ) is required for training [ 8 , 12 ] . However , labeled molecules only occupy an extremely small portion of the ∗Qi Liu and Chee-Kong Lee are corresponding authors . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . enormous chemical space since they can only be obtained from wet-lab experiments or quantum chemistry calculations , which are time-consuming and expensive . Moreover , directly training GNNs on small labeled molecule datasets in a supervised fashion is prone to over-fitting and the trained GNNs can hardly generalize to out-of-distribution data . Similar issues have also been encountered in natural language processing and computer vision . Recent advances in NLP and CV address them by self-supervised learning ( SSL ) where a model is first pre-trained on a large unlabeled dataset and then transferred to downstream tasks with limited labels [ 13 , 2 , 5 ] . For example , the pre-trained BERT language model [ 5 ] is able to learn expressive contextualized word representations through reconstructing the input text-next sentence and masked language predictions so that it can significantly improve the performance of downstream tasks . Inspired by these developments , various self-supervised pre-training methods of GNNs have been proposed [ 14 , 15 , 52 , 51 , 32 , 10 , 30 ] . Based on how the self-supervised tasks are constructed , these methods can be classified into two categories , contrastive methods and predictive methods . Contrastive methods force views from the same graph ( e.g. , sampling nodes and edges from graphs ) to become closer and push views from different graphs apart . On the other hand , the predictive methods construct prediction tasks by exploiting the intrinsic properties of data . For example , Hu et.al . [ 14 ] designed node-level pre-training tasks such as predicting the context of atoms and the attributes of masked atoms and bonds . [ 15 ] introduced an attributed graph reconstruction task where the generative model predicts the node attributes and edges to be generated at each step . However , we argue that existing self-supervised learning tasks on GNNs are sub-optimal since most of them fail to exploit the rich semantic information from graph motifs . Graph motifs can be defined as significant subgraph patterns that frequently occur [ 27 ] . Motifs usually contain semantic meanings and are indicative of the characteristics of the whole graph . For example , the hydroxide ( –OH ) functional group in small molecules typically implies higher water solubility . Therefore , it is vital to design motif-level self-supervised learning tasks which can benefit downstream tasks such as molecular property prediction . Designing motif-level self-supervised learning tasks brings unique challenges . First , existing motif mining techniques could not be directly utilized to derive expressive motifs for molecular graphs because they only rely on the discrete count of subgraph structures and overlook the chemical validity [ 1 , 18 ] . Second , most graph generation techniques generate graphs node-by-node [ 23 , 50 ] , which are not suitable for our task of motif generation . Finally , how to unify multi-level self-supervised pre-training tasks harmoniously brings a new challenge . One naive solution to do pre-training tasks sequentially may lead to catastrophic forgetting similar to continual learning [ 7 ] . To tackle the aforementioned challenges , we propose Motif-based Graph Self-Supervised Learning ( MGSSL ) and Multi-level self-supervised pre-training in this paper . Firstly , MGSSL introduces a novel motif generation task that empowers GNNs to capture the rich structural and semantic information from graph motifs . To derive semantically meaningful motifs and construct motif trees for molecular graphs , we leverage the BRICS algorithm [ 4 ] which is based on retrosynthesis from the chemistry domain . Two additional fragmentation rules are further introduced to reduce the redundancy of motif vocabulary . Secondly , a general motif-based generative pre-training framework is designed to generate molecular graphs motif-by-motif . The pre-trained model is required to make topology and attribute predictions at each step and two specific generation orders are implemented ( breadth-first and depth-first ) . Furthermore , to take the multi-scale regularities of molecules into consideration , we introduce Multi-level self-supervised pre-training for GNNs where the weights of different SSL tasks are adaptively adjusted by the Frank-Wolfe algorithm [ 16 ] . Finally , by pre-training GNNs on the ZINC dataset with our methods , the pre-trained GNNs outperforms all the state-of-the-art baselines on various downstream benchmark tasks , demonstrating the effectiveness of our design . The implementation is publicly available at https : //github.com/zaixizhang/MGSSL . 2 Preliminaries and Related Work . 2.1 Molecular Property Prediction . Prediction of molecular properties is a central research topic in physics , chemistry , and materials science [ 44 ] . Among the traditional methods , density functional theory ( DFT ) is the most popular one and plays a vital role in advancing the field [ 21 ] . However , DFT is very time-consuming and its complexity could be approximated as O ( N3 ) , where N denotes the number of particles . To find more efficient solutions for molecular property prediction , various machine learning methods have been leveraged such as kernel ridge regression , random forest , and convolutional neural networks [ 6 , 26 , 36 ] . To fully consider the internal spatial and distance information of atoms in molecules , many works [ 8 , 35 , 48 , 25 ] regard molecules as graphs and explore the graph convolutional network for property prediction . To better capture the interactions among atoms , [ 8 ] proposes a message passing framework and [ 20 , 48 ] extend this framework to model bond interactions . [ 25 ] builds a hierarchical GNN to capture multilevel interactions . Furthermore , [ 32 ] integrates GNNs into the Transformer-style architecture to deliver a more expressive model . 2.2 Preliminaries of Graph Neural Networks . Recent years have witnessed the success of Graph Neural Networks ( GNNs ) for modeling graph data [ 11 , 19 , 40 , 46 , 43 , 42 , 55 ] . Let G = ( V , E ) denotes a graph with node attributes Xv for v ∈ V and edge attributes euv for ( u , v ) ∈ E. GNNs leverage the graph connectivity as well as node and edge features to learn a representation vector ( i.e. , embedding ) hv for each node v ∈ G and a vector hG for the entire graph G. Generally , GNNs follows a message passing paradigm , in which representation of node v is iteratively updated by aggregating the representations of v ’ s neighboring nodes and edges . Specifically , there are two basic operators for GNNs : AGGREGATE ( · ) and COMBINE ( · ) . AGGREGATE ( · ) extracts the neighboring information of node v ; COMBINE ( · ) serves as the aggregation function of the neighborhood information . After k iterations of aggregation , v ’ s representation captures the structural information within its k-hop network neighborhood . Formally , the k-th layer of a GNN is : h ( k ) v = COMBINE ( k ) ( h ( k−1 ) v , AGGREGATE ( k ) ( { ( h ( k−1 ) v , h ( k−1 ) u , euv ) : u ∈ N ( v ) } ) ) , ( 1 ) where h ( k ) v denotes the representation of node v at the k-th layer , euv is the feature vector of edge between u and v and andN ( v ) represents the neighborhood set of node v. h ( 0 ) v is initialized with Xv . Furthermore , to obtain the representation of the entire graph hG , READOUT functions are designed to pool node representations at the final iteration K , hG = READOUT ( { h ( K ) v | v ∈ G } ) . ( 2 ) READOUT is a permutation-invariant function , such as averaging , sum , max or more sophisticated graph-level pooling functions [ 49 , 53 ] . 2.3 Self-supervised Learning of Graphs . Graph Self-supervised learning aims to learn the intermediate representations of unlabeled graph data that are useful for unknown downstream tasks . [ 14 , 15 , 52 , 51 , 32 , 10 , 30 , 39 , 28 , 41 ] . Traditional graph embedding methods [ 9 , 29 ] define different kinds of graph proximities , i.e. , the vertex proximity , as the self-supervised objective to learn vertex representations . Furthermore , [ 39 , 28 , 41 ] proposed to use the mutual information maximization as the optimization objective for GNNs . Recently , more self-supervised tasks for GNNs have been proposed [ 14 , 15 , 52 , 51 , 32 , 10 , 54 , 30 ] . Based on how the self-supervised tasks are constructed , these models can be classified into two categories , namely contrastive models and predictive models . Contrastive models try to generate informative and diverse views from data instances and perform node-to-context [ 14 ] , node-to-graph [ 39 ] or motif-to-graph [ 54 ] contrastive learning . On the other hand , predictive models are trained in a supervised fashion , where the labels are generated based on certain properties of the input graph data , i.e. , node attributes [ 32 ] , or by selecting certain parts of the graph [ 15 , 32 ] . However , few works try to leverage the information of graph motifs for graph self-supervised learning . Grover [ 32 ] use traditional softwares to extract motifs and treat them as classification labels . MICROGraph [ 54 ] exploits graph motifs for motif-graph contrastive learning . Unfortunately , both methods fail to consider the topology information of motifs . 3 Motif-based Graph Self-supervised Learning . In this section , we introduce the framework of motif-based graph self-supervised learning ( Figure 1 ) . Generally , the framework consists of three parts : chemistry-inspired molecule fragmentation , motif generation and multi-level self-supervised pre-training for GNNs . 3.1 Chemistry-inspired Molecule Fragmentation . Given a dataset of molecules , the first step of our method is to decompose molecules into several fragments/motifs . Based on these fragments , a molecule graph can be converted into a motif tree structure where each node represents a motif and the edges denotes the relative spatial relationships among motifs . We choose to represent the structure of motifs as a tree because it facilitates the motif generation task . Formally , given a molecule graph G = ( V , E ) , a motif tree T ( G ) = ( V , E , X ) is a connected labeled tree whose node set is V = { M1 , ... , Mn } and edge set is E . X refers to the induced motif vocabulary . Each motif Mi = ( Vi , Ei ) is a subgraph of G. There are many ways to fragment a given graph while the designed molecule fragmentation method should achieve the following goals : 1 ) In a motif tree T ( G ) , the union of all motifs Mi should equals G. Formally , ⋃ i Vi = V and ⋃ iEi ⋃ E = E. 2 ) In a motif tree T ( G ) , the motifs should have no intersections . That is Mi ∩Mj = ∅ . 3 ) The induced motifs should capture semantic meanings , e.g. , similar to meaningful functional groups in the chemistry domain . 4 ) The occurrence of motifs in the dataset should be frequent enough so that the pre-trained GNNs can learn semantic information of motifs that can be generalized to downstream tasks . After the molecule fragmentation , the motif vocabulary of the molecule dataset should have a moderate size . In Figure 2 , we show the overview of molecule fragmentation . Generally , there are three procedures , BRICS fragmentation , further decomposition , and motif tree construction . A motif vocabulary can be built via preprocessing the whole molecule dataset following the molecule fragmentation procedures . To fragment molecule graphs and construct motif trees , we firstly use the Breaking of Retrosynthetically Interesting Chemical Substructures ( BRICS ) algorithm [ 4 ] that leverages the domain knowledge from chemistry . BRICS defines 16 rules and breaks strategic bonds in a molecule that match a set of chemical reactions . `` Dummy '' atoms are attached to each end of the cleavage sites , marking the location where two fragments can join together . BRICS cleavage rules are designed to retain molecular components with valuable structural and functional content , e.g . aromatic rings . However , we find BRICS alone can not generate desired motifs for molecule graphs . This is because BRICS only breaks bonds based on a limited set of chemical reactions and tends to generate several large fragments for a molecule . Moreover , due to the combination explosion of graph structure , we find BRICS produces many motifs that are variations of the same underlying structure ( e.g. , Furan ring with different combinations of Halogen atoms ) . The motif vocabulary is large ( more than 100k unique fragments ) while most of these motifs appear less than 5 times in the whole dataset . To tackle the problems above , we introduce a post-processing procedure to BRICS . To alleviate the combination explosion , we define two rules operating on the output fragments of BRICS : ( 1 ) Break the bond where one end atom is in a ring while the other end not . ( 2 ) Select non-ring atoms with three or more neighboring atoms as new motifs and break the neighboring bonds . The first rule reduces the number of ring variants and the second rule break the side chains . Experiments show that these rules effectively reduce the size of motif vocabulary and improve the occurrence frequency of motifs .

Summary: This paper presents a new self-supervised pre-training framework for Graph Neural Networks (GNNs) called Motif-based Graph Self-supervised Learning (MGSSL) to predict molecular properties. First, the method uses BRICs, a molecule fragmentation method, to create a Graph Motif vocabulary. Next, it performs pre-training using a model with three components -  1) a motif generation module that uses GRUs to recursively generate motifs given the previously generated motifs. 2)  a motif labeling module that uses dense layers followed by a softmax to predict the label of the motif from the vocabulary created in the previous step. 3) a GNN module that generates node and edge embeddings which are finally used to predict node and edge attributes of a molecule graph. MGSSL is trained as a multi-level model (with atom and motif layers) and it minimizes a hybrid loss consisting of Motif, Atomic, and Bond prediction losses, with weights adjusted by the Frank-Wolfe Algorithm. The results report improved performance for molecular property predictions over existing pre-training frameworks 7/8 datasets.