Property-Aware Relation Networks for Few-Shot Molecular Property Prediction

1 Introduction . Drug discovery is an important biomedical task , which targets at finding new potential medical compounds with desired properties such as better absorption , distribution , metabolism , and excretion ( ADME ) , low toxicity and active pharmacological activity [ 1 , 2 , 3 ] . It is recorded that drug discovery takes more than 2 billion and at least 10 years in average while the clinical success rate is around 10 % [ 4 , 5 , 6 ] . To speed up this process , quantitative structure property/activity relationship ( QSPR/QSAR ) modeling uses machine learning methods to establish the connection between molecular structure and particular properties [ 7 ] . It usually consists of two components : a molecular encoder which encodes molecular structure as a fixed-length molecular representation , and a predictor which estimates the activity of a certain property based on the molecular representation . Predictive models can be leveraged in virtual screening to discover potential molecules more efficiently [ 8 ] . However , molecular property prediction is essentially a few-shot problem , which makes it hard to solve . Only a small amount of candidate molecules can pass virtual screening to be evaluated in the lead optimization stage of drug discovery [ 9 ] . After a series of wet-lab experiments , most candidates ⇤Equal contribution . A. Abuduweili did his work during internship at Baidu Research . †Correspondence to . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . eventually fail to be a potential drug due to the lack of any desired properties [ 7 ] . These together result in a limited number of labeled data [ 10 ] . Few-shot learning ( FSL ) [ 11 , 12 ] methods target at generalizing from a limited number of labeled data . Recently , they have also been introduced into molecular property prediction [ 3 , 8 ] . These methods attempt to learn a predictor from a set of property prediction tasks and generalize to predict new properties given a few labeled molecules . As molecules can be naturally represented as graphs , graphbased molecular representation learning methods use graph neural networks ( GNNs ) [ 13 , 14 ] to obtain graph-level representation as the molecular embedding . Specifically , the pioneering IterRefLSTM [ 3 ] adopts GNN as the molecular encoder and adapts a classic FSL method [ 15 ] proposed for image classification to handle few-shot molecular prediction tasks . The recent Meta-MGNN [ 8 ] leverages a GNN pretrained from large-scale self-supervised tasks as molecular encoder and introduces additional self-supervised tasks such as bond reconstruction and atom type prediction to be jointly optimized with the molecular property prediction tasks . However , aforementioned methods neglect two key facts in molecular property prediction . The first fact is that different molecular properties are attributed to different molecular substructures as found by previous QSPR studies [ 16 , 17 , 18 ] . However , IterRefLSTM and Meta-MGNN use graph-based molecular encoder to encode molecules regardless of target properties whose relevant substructures are quite different . The second fact is that the relationship among molecules also vary w.r.t . the target property . This can be commonly observed in benchmark molecular property prediction datasets . As shown in Figure 1 , Mol-1 and Mol-4 from the Tox21 dataset [ 19 ] have the same activity in SR-HSE task while acting differently in SR-MMP task . However , existing works fail to leverage such relation graph among molecules . To handle these problems , we propose Property-Aware Relation networks ( PAR ) which is compatible with existing graph-based molecular encoders , and is further equipped with the ability to obtain property-aware molecular embeddings and model molecular relation graph adaptively . Specifically , our contribution can be summarized as follows : • We propose a property-aware embedding function which co-adapts each molecular embedding with respect to context information of the task and further projects it to a substructure-aware space w.r.t . the target property . • We propose an adaptive relation graph learning module to jointly estimate molecular relation graph and refine molecular embeddings w.r.t . the target property , such that the limited labels can be effectively propagated among similar molecules . • We propose a meta-learning strategy to selectively update parameters within each task , which is particularly helpful to separately capture the generic knowledge shared across different tasks and those specific to each property prediction task . • We conduct extensive empirical studies on real molecular property prediction datasets . Results show that PAR consistently outperforms the others . Further model analysis shows PAR can obtain property-aware molecular embeddings and model molecular relation graph properly . Notation . In the sequel , we denote vectors by lowercase boldface , matrices by uppercase boldface , and sets by uppercase calligraphic font . For a vector x , [ x ] i denotes the ith element of x . For a matrix X , [ X ] i : denotes the vector on its ith row , [ X ] ij denotes the ( i , j ) th element of X . The superscript ( · ) > denotes the matrix transpose . 2 Review : Graph Neural Networks ( GNNs ) . A graph neural network ( GNN ) can learn expressive node/graph representation from the topological structure and associated features of a graph via neighborhood aggregation [ 13 , 20 , 21 ] . Consider a graph G = { V , E } with node feature h ( 0 ) v for each node v 2 V and edge feature bvu for each edge evu 2 E between nodes v , u . At the lth layer , GNN updates the node embedding h ( l ) v of node v as : h ( l ) v = UPDATE ( l ) ⇣ h ( l 1 ) v , AGGREGATE ( l ) ⇣ { ( h ( l 1 ) v , h ( l 1 ) u , bvu ) |u 2 N ( v ) } ⌘⌘ , ( 1 ) where N ( v ) is a set of neighbors of v. After L iterations of aggregation , the graph-level representation g for G is obtained as g = READOUT ⇣ { h ( L ) v |v 2 V } ⌘ , ( 2 ) where READOUT ( · ) function aggregates all node embeddings into the graph embedding [ 22 ] . Our paper is related to GNN in two aspects : ( i ) use graph-based molecular encoder to obtain molecular representation , and ( ii ) conduct graph structure learning to model relation graph among molecules . Graph-based Molecular Representation Learning . Representing molecules properly as fixedlength vectors is vital to the success of downstream biomedical applications [ 23 ] . Recently , graphbased molecular representation learning methods are popularly used and obtain state-of-the-art performance . A molecule xi is represented as an undirected graph Gi = { Vi , Ei } , where each node v 2 Vi represents an atom with feature h ( 0 ) v 2 Rd n and each edge evu 2 Ei represents the bond between two nodes v , u with feature bvu 2 Rd e . Graph-based molecular representation learning methods use GNNs to obtain graph-level representation gi as molecular embedding . Examples include graph convolutional networks ( GCN ) [ 24 ] , graph attention networks ( GAT ) [ 25 ] , message passing neural networks ( MPNN ) [ 20 ] , graph isomorphism network ( GIN ) [ 22 ] , pretrained GNN ( Pre-GNN ) [ 26 ] and GROVER [ 9 ] . Existing two works in few-shot molecular property prediction both use graph-based molecular encoder to obtain molecular embeddings : IterRefLSTM [ 3 ] uses GCN while Meta-MGNN [ 8 ] uses Pre-GNN . Using these graph-based molecular encoders can not discover molecular substructures corresponding to the target property . There exist GNNs which handle subgraphs [ 27 , 28 , 29 ] , which are usually predefined or simply K-hop neighborhood . While discovering and enumerating molecular substructures is extremely hard even for domain experts [ 17 , 30 ] . In this paper , we first obtain molecular embeddings using graph-based molecular encoders . We further learn to extract relevant substructure embeddings w.r.t . the target property upon these generic molecular embeddings , which is more effective and improves the performance . Graph Structure Learning . As the provided graphs may not be optimal , a number of graph structure learning methods target at jointly learning graph structure and node embeddings [ 31 , 32 , 33 ] . In general , they iterate over two procedures : ( i ) estimate adjacency matrix ( i.e. , refining neighborhood u 2 N ( v ) ) which encodes graph structure from the current node embeddings ; and ( ii ) apply a GNN on this updated graph to obtain new node embeddings . There exist some FSL methods [ 34 , 35 , 36 , 37 , 38 ] which learn to construct fully connected relation graph among images in a N -way K-shot few-shot image classification task . Their methods can not work for the 2-way K-shot property prediction tasks where choosing a wrong neighbor in the different class will heavily deteriorate the quality of molecular embeddings . We share the same spirit of learning relation graph , and further design several regularizations to encourage our adaptive property-aware relation graph learning module to select correct neighbors . 3 Proposed Method . In this section , we present the details of PAR , whose overall architecture is shown in Figure 2 . Considering few-shot molecular property prediction problem , we first use a specially designed embedding function to obtain property-aware molecular embedding for each molecule , and then adaptively learn relation graph among molecules which allows effective propagation of the limited labels . Finally , we describe our meta-learning strategy to train PAR . 3.1 Problem Definition . Following the problem definition adopted by IterRefLSTM [ 3 ] and Meta-MGNN [ 8 ] , the target is to learn a predictor from a set of few-shot molecular property prediction tasks { T⌧ } Nt⌧=1 and generalize to predict new properties given a few labeled molecules . The ⌧ th task T⌧ predicts whether a molecule x⌧ , i with index i is active ( y⌧ , i = 1 ) or inactive ( y⌧ , i = 0 ) on a target property , provided with a small number of K labeled samples per class . This T⌧ is then formulated as a 2-way K-shot classification task with a support set S⌧ = { ( x⌧ , i , y⌧ , i ) } 2Ki=1 containing the 2K labeled samples and a query set Q⌧ = { ( x⌧ , j , y⌧ , j ) } Nq⌧ j=1 containing N q ⌧ unlabeled samples to be classified . 3.2 Property-aware Molecular Embedding . As different molecular properties are attributed to different molecule substructures , we design a property-aware embedding function to transform the generic molecular embeddings to substructureaware space relevant to the target property . As introduced in Section 2 , graph-based molecular encoders can obtain good molecular embeddings . By learning from large-scale tasks , they can capture generic information shared by molecules [ 26 , 9 ] . Thus , we first use a graph-based molecular encoder such as GIN [ 22 ] and Pre-GNN [ 26 ] to extract a molecular embedding g⌧ , i 2 Rd g of length dg for each x⌧ , i . The parameter of this graph-based molecular encoder is denoted as Wg . However , existing graph-based molecular encoders can not capture property-aware substructures . Especially when learning across tasks , a molecule can be evaluated for multiple properties . This leads to a one-to-many relationship between a molecule and properties , which makes few-shot molecular property prediction particularly hard . Thus , we are motivated to implicitly capture substructures in the embedding space w.r.t . the target property of T⌧ . Let cc⌧ denote the class prototype for class c 2 { 0 , 1 } , which is computed as c c ⌧ = 1/|S c ⌧ | X ( x⌧ , i , y⌧ , i ) 2Sc⌧ g⌧ , i , ( 3 ) where Sc⌧ = { ( x⌧ , i , y⌧ , i ) | ( x⌧ , i , y⌧ , i ) 2 S⌧ and y⌧ , i = c } . We take these class prototypes as the context information of T⌧ , and further encode them as b⌧ , i = ⇥ softmax ( C⌧ , iC > ⌧ , i/ p dg ) C⌧ , i ⇤ 1 : with C > ⌧ , i = [ g⌧ , i , c 0 ⌧ , c 1 ⌧ ] 2 Rd g⇥3 , ( 4 ) where [ · ] j : extracts the jth row vector which corresponds to x⌧ , i . Here b⌧ , i is computed using scaled dot-product self-attention [ 39 ] , such that each g⌧ , i can be compared with class prototypes in a dimensional wise manner . The property-aware molecular embedding p⌧ , i is then obtained as p⌧ , i = MLPWp ( concat [ g⌧ , i , b⌧ , i ] ) . ( 5 ) MLPWp denotes the multilayer perceptron ( MLP ) parameterized by Wp , which is used to find a lower dimensional space which encodes substructures that are more relevant to the target property of T⌧ . This contextualized p⌧ , i is property-aware which can be more predictive of the target property .

Summary: This paper proposes a property-aware adaptive relation networks for few shot molecular property prediction.  They propose a graph based molecule encoder which is trained from large scale tasks to capture a generic representation of the molecules and a property aware embedding function which adapts the generic molecule embedding to be property-aware. Next, an adaptive relation graph learning module is proposed to capture the relationship between similar molecules with respect to a target property. The adaptive relation graph learning module is based on GNNs. A new training strategy, where only the parameters of property aware embedding function and the classifier module is tuned, is also proposed. The authors argue that it useful to separately capture generic molecule representations and those that are specific to the particular task.