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: The authors look at molecular property prediction from a few shot learning perspective. They propose an approach, closely following Meta-MGNN, to adapt part of a (optionally pre-trained) graph embedding network in MAML-style per task. A number of well-motivated novelties are introduced: (a) adapt part of the GNN encoder, called property-aware transformation (vs just task-specific attention in Meta-MGNN). (b) explicitly constructing the similarity graph between a test sample and the support set (labeled samples) and add a sparsity loss for the adjacency matrix during training. The results, focusing on 1-shot and 10-shot prediction, convincingly show gains in this regime, and extensive ablations and baselines are presented.