Towards Interpretable Molecular Graph Representation Learning

1 INTRODUCTION . Following the recent rise of deep learning for image and speech processing , there has been great interest in generalizing convolutional neural networks to arbitrary graph-structured data ( Gilmer et al. , 2017 ; Henaff et al. , 2015 ; Xu et al. , 2018 ) . To this end , graph neural networks ( GNN ) falling into either spectral-based or spatial-based approaches have been proposed . Spectral methods define the graph convolution ( GC ) as a filtering operator of the graph signal ( Defferrard et al. , 2016 ) , while spatial methods define the GC as a message passing and aggregation across nodes ( Henaff et al. , 2015 ; Xu et al. , 2018 ; Jin et al. , 2018 ) . In drug discovery , GNNs have been very successful across several molecular graph classification and generation tasks . In particular , they outperform predetermined molecular fingerprints and string-based approaches for molecular property prediction and the de novo generation of drug-like compounds ( Jin et al. , 2018 ; Li et al. , 2018b ) . However , the node feature update performed by GNNs introduces important limitations . For instance , experimental results indicate a performance decrease for deeper GNNs due to the signal smoothing effect of each GC layer ( Li et al. , 2018a ) . This limits the network ’ s depth and restricts the receptive field of the vertices in the graph to a few-hop neighborhood , which is insufficient to properly capture local structures , relationships between nodes , and subgraph importance in sparse graphs such as molecules . For example , at least three consecutive GC layers are needed for atoms at the opposite side of a benzene ring to exchange information . This issue is exacerbated by the single global pooling step performed at the end of most GNNs that ignores any hierarchical structure within the graph . To cope with these limitations , graph coarsening ( pooling ) methods have been proposed to reduce graph size and enable long-distance interaction between nodes . The earliest proposed methods relied solely on deterministic clustering of the graphs , making them non-differentiable and task-independent ( Jin et al. , 2018 ; Dafna and Guestrin , 2009 ; von Luxburg , 2007 ; Ma et al. , 2019 ) . In contrast , more recent methods use node features but , as we will show , are unable to preserve the graph structures after pooling ( Ying et al. , 2018 ; Gao and Ji , 2018 ) , limiting their interpretability . Borrowing from theory in graph signal processing , we propose LaPool ( Laplacian Pooling ) , a differentiable pooling method that takes into account both the graph structure and its node features . LaPool performs a dynamic and hierarchical segmentation of graphs by selecting a set of centroid nodes as cluster representatives ( centroids ) using the graph Laplacian , then learns a sparse assignment of the remaining nodes ( followers ) into these clusters using an attention mechanism . The primary contributions of this paper are summarized below : • We propose a novel and differentiable pooling module ( LaPool ) that can be incorporated into existing GNNs to yield more expressive networks for molecular data . • We propose a graph structure understanding dataset for benchmarking GNNs that is based on molecular substructure prediction . • We show that LaPool outperforms recently proposed graph pooling layers on both discrimi- native and generative molecular graph benchmarks , while also remaining competitive on other graph benchmarks . • We highlight the improved interpretability achieved by LaPool using both qualitative and quantitative assessments . We argue that the enhanced performance and interpretability achieved by LaPool can improve our understanding of molecular structure-activity relationships , and therefore has important applications in drug discovery . 2 RELATED WORK . In this section , we first introduce related work on graph pooling , then provide an overview of techniques used in computational drug discovery to put our work into context . As our focus herein is on graph pooling , we refer readers to ( Wu et al. , 2019 ) for an overview of recent progress in GNNs . In traditional GNN architectures , global sum/average/max-pooling layers have been used to aggregate node embeddings into a graph-level representation . Recently , more sophisticated methods have been proposed . For example , Li et al . ( 2015 ) uses a gated mechanism , Zhang et al . ( 2018 ) proposed SortPool which sorts node features before feeding them into a 1D convolution , while in ( Gilmer et al. , 2017 ) node feature averaging was substituted by a Set2Set architecture . Although these new global aggregation methods have been shown to outperform standard global pooling , they completely overlook the rich structural information on graphs which has been shown as necessary for building effective GNN models ( Ying et al. , 2018 ; Ma et al. , 2019 ) . Consequently , hierarchical graph pooling methods have been proposed . They act by reducing graph size and increasing node receptive fields without increasing network depth . However , in contrast to the regular structure of images , graphs are irregular and complex , making it challenging to properly pool nodes together . Certain hierarchical graph pooling methods therefore rely on deterministic and non-differentiable clustering to segment the graph ( Defferrard et al. , 2016 ; Jin et al. , 2018 ) . More recently , differentiable hierarchical graph pooling layers have been proposed . Ying et al . ( 2018 ) proposed DiffPool , a pooling layer that performs a similarity-based node clustering using a soft affinity matrix learned by a GNN . Likewise , Graph U-Net was proposed in ( Gao and Ji , 2018 ) as a sampling method that retains and propagates only the top-k nodes at each pooling step based on a learned importance value . In computer-aided drug discovery , methods such as virtual screening and de novo drug design serve as efficient complements to physical high-throughput screening of large molecular libraries . For example , virtual screening , which aims to accurately predict molecular properties directly from molecular structure , can play an important role in rapidly triaging promising compounds early in drug discovery ( Subramaniam et al. , 2008 ) . Importantly , data-driven virtual screening approaches that leverage advances in deep learning , rather than pre-determined features such as molecular fingerprints ( Rogers and Hahn , 2010 ) and SMILES string representations , have been shown to dramatically improve prediction accuracy ( Kearnes et al. , 2016 ; Wu et al. , 2018 ) . Similarly , advances in generative models have enabled the application of deep generative techniques such as VAE ( Kingma and Welling , 2013 ) and GAN ( Goodfellow et al. , 2014 ) to the de novo design of drug-like molecules . The first molecular generative models ( e.g . Grammar-VAE ( Kusner et al. , 2017 ) ) resorted to generating string representations of molecules ( via SMILES ) , which resulted in a high proportion of invalid structures due to the complex syntax of SMILES . Graph generative models have since been developed ( e.g . JT-VAE ( Jin et al. , 2018 ) , GraphVAE ( Simonovsky and Komodakis ) , MolGAN ( De Cao and Kipf , 2018 ) , MolMP ( Li et al. , 2018b ) , etc . ) and have been shown to improve the validity and novelty of generated molecules . In addition , these methods allow conditional molecule generation via Bayesian optimization or reinforcement learning ( Jin et al. , 2018 ; Olivecrona et al. , 2017 ; Assouel et al. , 2018 ; Li et al. , 2018d ; You et al. , 2018a ) . In this work , we are mainly interested in the impact of molecular representation on generative performance as opposed to the optimization procedure itself . 3 GRAPH LAPLACIAN POOLING . A reliable pooling operator should maintain the overall structure and connectivity of a graph . LaPool achieves this by taking into account the local structure defined by the neighborhood of each node . As shown in Figure 1 , the method uses a standard GC layer with a centroid selection and a follower selection step . First , the centroids of the graph are selected based on the local signal variation ( see Section 3.2 ) . Next , LaPool learns an affinity matrix C using a distance normalized attention mechanism to assign all nodes of the graph to the centroids ( see Section 3.3 ) . Finally , the affinity matrix is used to coarsen the graph . These steps are detailed below . 3.1 PRELIMINARIES . Notation Let G = 〈V , A , X〉 be an undirected graph , where V = { v1 , . . . vn } is its vertex set , A ∈ { 0 , 1 } n×n denotes its adjacency matrix , and X = [ x1 , . . .xn ] ᵀ ∈ Rn×d is the node feature matrix with each node vi having d-dimensional feature xi . X can also be viewed as a d-dimensional signal on G ( Shuman et al. , 2012 ) . Without loss of generality we may assume a fixed ordering of the nodes that is respected in V , A , and X . The neighborhood of radius h ( or h-hop ) neighborhood of a node vi ∈ V is the set of nodes separated from vi by a path of length at most h and is denoted by N h ( vi ) . For simplicity , we will use N ( vi ) to refer to the set of nodes adjacent to vi . Graph Signal For any graph G , its graph Laplacian matrix L is defined as L = D −A , where D is a diagonal matrix with Di , i being the degree of node vi . The graph Laplacian is a difference operator and can be used to define the smoothness s ( X ) ( the extent at which the signal changes between connected nodes ) of a signal X on G. For 1-dimensional signal f = [ f1 , . . . , fn ] : s ( f ) = ( fᵀLf ) = 1 2 n∑ i , j Ai , j ( fi − fj ) 2 ( 1 ) Graph Neural Networks We consider GNNs that act in the graph spatial domain as message passing ( Gilmer et al. , 2017 ) . We focus on the Graph Isomorphism Network ( GIN ) ( Xu et al. , 2018 ) , which uses a SUM-aggregator on messages received by each node to achieve a better understanding of the graph structure : xli = M l Θ x ( l−1 ) i + ∑ vj∈N ( vi ) A ( l−1 ) i , j x ( l−1 ) j ( 2 ) where MlΘ is a neural network with trainable parameters Θ , xi is the feature vector for node vi , vj ∈ N ( vi ) are the neighbors of vi and l is the layer number . Notice the term Ai , j that takes into account the edge weight between nodes vi and vj when A is not binary . In this work , we focus mainly on molecular graphs in supervised settings where , given a molecule m and its corresponding molecular graph Gm , we aim to predict some properties of m. Molecular graphs present two particularities : ( 1 ) they are often sparse and ( 2 ) there is no regularity in the graph signal ( non-smooth variation ) as adjacent nodes tend not to have similar features .

Summary: The paper introduces a new pooling approach "Laplacian pooling" for graph neural networks, which the authors claim is able to better preserve information about the local structure, and to provide interpretability.  Namely, the pooling approach is based on finding centroids (nodes having high signal variation compared to their neighbors, via graph-Laplacian) and assigning other nodes to be "followers" based on a soft-attention mechanism. The authors add these new pooling layers to existing GNN architectures and show improved performance on problems of classification and generative modeling of molecular graphs. The paper also extends CNN interpretability techniques (integrated gradients) to GNNs.