Year: 2019

Abstract:
Recent work in graph neural networks (GNNs) has led to improvements in molecular activity and property prediction tasks. Unfortunately, GNNs often fail to capture the relative importance of interactions between molecular substructures, in part due to the absence of efficient intermediate pooling steps. To address these issues, we propose LaPool (Laplacian Pooling), a novel, data-driven, and interpretable hierarchical graph pooling method that takes into account both node features and graph structure to improve molecular understanding. We benchmark LaPool and show that it not only outperforms recent GNNs on molecular graph understanding and prediction tasks but also remains highly competitive on other graph types. We then demonstrate the improved interpretability achieved with LaPool using both qualitative and quantitative assessments, highlighting its potential applications in drug discovery.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 GRAPH LAPLACIAN POOLING
text: 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.

heading: 3.1 PRELIMINARIES
text: 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.

heading: 3.2 GRAPH DOWNSAMPLING
text: This section details how LaPool downsamples the original graph by selecting a set VC of nodes as centroids after l consecutive GC layers.
Centroid Selection For any given vertex vi, we can define a local measure si of intensity of signal variation around vi. As si measures how different the signal residing at a node vi is from its neighbors, we are interested in the set VC of nodes that have an intensity of signal variation si greater than their neighborhood. In this work, we use a definition of local signal variation similar to the local normalized neighboring signal variation described in (Chen et al., 2015b), with the only difference being the absence of degree normalization:
si = ∥∥∑ j∈N (vi) Ai,j(xi − xj) ∥∥ 2 , S = [ s1, . . . sn ]ᵀ = ‖LX‖2,Rd VC = topV (LS, k) (3)
where topV (L · S, k) corresponds to the top k nodes with the greatest intensity of signal variation, and where || · ||2,Rd corresponds to taking the vector norms over the d-dimensional rows of LX . Instead of using the direct neighbors, one can also generalize the computation of S in equation 3 to an h−hop neighborhood by taking powers of the Laplacian. Observe that the GC layers preceding each pooling step perform a smoothing of the graph signal and thus act as a low-pass filter. Eq. (3) emphasizes instead the high variation regions, resulting overall in a filtering of X that attenuates low and high-frequency noise, yet retains the important signal information. The intuition of using the Laplacian maxima for selecting the centroids is that a smooth signal can be very well approximated using a linear interpolation between its local maxima and minima. This is in contrast with most approaches in GSP that use the lower frequencies for signal conservation but requires the signal to be k-bandlimited (Ma et al., 2019; Chen et al., 2015c;a). For a 1D signal, LaPool selects points, usually near the maxima/minima, where the derivative changes the most and is hardest to interpolate linearly. For molecular graphs, this often corresponds to sampling a subset of nodes critical for reconstructing the original molecule.
Dynamic Selection of the Centroids The method presented in Eq. (3) implies the selection of k centroids. Because the optimal value of k can be graph-dependent and might result in densely located centroids, we explore alternative in which we dynamically select the nodes with signal variation si greater than its neighbors sj :
VC = {vi ∈ V | ∀ vj , si −Aijsj > 0} (4)

heading: 3.3 LEARNING THE NODE-TO-CLUSTER ASSIGNMENT MATRIX
text: Once the set VC of centroid nodes is determined, we compute a mapping of the remaining “follower” nodes VF = V \ VC into the new clusters formed by the nodes in VC . This mapping gives the cluster assignment C = [ c1, ...cn ]ᵀ ∈ [0, 1]n×m s.t. ∀i : 1cᵀi = 1, where each row ci corresponds to the affinity of node vi towards each of the m clusters in VC .
LetX(l) be the node embedding matrix at an arbitrary layer andX(l)C the embedding of the “centroids”. We compute C using a soft-attention mechanism (Graves et al., 2014) measured by the cosine similarity between X(l) and X(l)C :
ci = δi,j if vi ∈ VCsparsemax(βi xi(l)·X(l)C‖xi(l)‖‖X(l)C ‖ ) otherwise (5)
where δi,j is the Kronecker delta and sparsemax (Laha et al., 2016; Martins and Astudillo, 2016), is an alternative to the softmax operator defined by:
sparsemax(z) = arg min p∈∆K−1 ‖p − z‖2 which corresponds to the euclidean projection of z onto the probability simplex ∆K−1 = { p ∈ RK |1ᵀp = 1,p ≥ 0 } . The sparsemax operator ensures the sparsity of the attention coefficients and encourages the assignment of each node to a single centroid. It further alleviates the need for entropy minimization as done in DiffPool.
Eq. (5) also prevents the selected centroid nodes from being assigned to other clusters. Moreover, notice the term βi that regularizes the value of the attention for each node. We can define βi = 1di,VC
, where di,VC is the shortest path distance between each node vi ∈ VF and centroids in VC . Although this regularization incurs a cost O(|V |2|VC |), it will strengthen the affinity to closer centroids and ensure the connectivity of the resulting pooled graph. Note that this regularization is considered as a hyperparameter of the layer and can be turned off, or alternatively, the mapping can be restricted to centroids within a fixed h−hop neighborhood of each node.
Finally, after Cl is computed at layer l, the coarsened graph G(l+1) = 〈V (l+1), A(l+1), X(l+1)〉 is computed using Eq. (6), as in (Ying et al., 2018). In these equations, MΨ is a neural network with trainable parameters Ψ that is used to update the embedding of nodes in G(l+1) after the mapping.
A(l+1) = C(l) ᵀ A(l)C(l) ∈ R|V (l) C |×|V (l) C |, X(l+1) = MΨ ( C(l) ᵀ X(l) ) (6)
This process can be repeated by feeding the new graph G(l+1) into another GNN layer.

heading: 3.4 PROPERTIES OF THE LAPOOL METHOD
text: Preservation of Structural Information In addition to identifying graph nodes as centroids for pooling in a data-driven way, LaPool retains the feature content of the other nodes in a graph via the soft-assignment of followers to their centroids.
Substructure Identification By construction, the soft assignment of nodes to centroids clusters existing substructures of the graph together, thus identifying important subgraphs according to the classification task. By controlling for differences in signal variations within neighborhoods, we encourage these clusters to be spread out across different areas of the graph.
Dynamic Cluster Dimension As discussed in section 3.2, LaPool offers the unique flexibility of determining the clustering dynamically, when training graphs sequentially, or statically when performing batch training.
Permutation Invariance It is trivial to show that LaPool is permutation invariant as long as the GNN used as its basis is permutation invariant, since both the graph downsampling (Eq. 3,4) and the node mapping (Eq. 5,6) are not affected by any permutation on the vertex set.
Emphasizing the Strong Features Similar to how most CNNs implement a max-pooling layer to emphasize the strong features, LaPool does so by selecting the nodes with high signal as centroids. For molecular graphs, the centroids are biased towards high degree nodes and atoms different than their neighbors (e.g. a Nitrogen in a Carbon ring).

heading: 4 RESULTS AND DISCUSSION
text: A fundamental objective of LaPool is to learn an interpretable representation of sparse graphs, notably molecular substructures. We argue that this is an essential step towards shedding light upon the decision process within neural networks and ultimately increasing their utility in the design of new drug-like molecules. This implies that GNNs should be able to identify semantically important substructure components from molecular graphs, and eventually reconstruct these graphs from such components. This stems from the intuition that molecular validity and functional properties derive more from chemical fragments than individual atoms.
Our experimental results thus aim to empirically demonstrate the following properties of LaPool, as benchmarked against current state-of-the-art pooling models and the Graph Isomorphism Network:
• LaPool’s consideration of semantically important information such as node distance translates to improved performance on molecular understanding and molecular activity prediction tasks. • Visualization of LaPool’s behaviour at the pooling layer demonstrates its ability to identify
coherent and meaningful molecular substructures. • The hierarchical representation enforced by LaPool, which preserves the original graph
structure improves model interpretability.
• Learning meaningful substructures can be leveraged to construct a generative model which leads to more realistic and feasible molecules.
Throughout our experiments, we use the same architecture for all models to ensure an even comparison across all pooling layers: 2 layer GNNs with 128 filters each before the optional pooling layer, followed by 2 GNNs with 64 filters each and two fully connected layers including the output layer. Detailed information on architectural tuning and pooling-specific hyper-parameter search are provided in Supplemental Section A.

heading: 4.1 BENCHMARK ON MOLECULAR GRAPH UNDERSTANDING
text: DiffPool and Graph U-Net models have been shown to outperform standard graph convolution networks on several graph benchmark datasets (Ying et al., 2018; Gao and Ji, 2018). Although not explicitly stated, both methods are most effective when the graph signal is smooth. In such cases where adjacent nodes tend to be similar, the DiffPool procedure will cluster together nodes in the same neighborhood, maintaining the overall graph structure, while Graph U-net will select nodes in the same neighborhood and will not create isolated components that no longer exchange information. On molecular graphs, however, the graph signal is rarely smooth. Therefore, we expect these two methods to be less effective at identifying the important molecular substructures,given that they do not explicitly consider structural relationships. We demonstrate this empirically by extracting known molecular substructure information from publicly available molecular datasets and evaluating performance in identifying these structures. We use a subset of approximately 17,000 molecules extracted from the ChEMBL database (Li et al., 2018c) and benchmark all methods on different types of substructures to verify the robustness of the comparison.
As shown in Table 1, the capture of structural relationships translates into superior performance of LaPool, as measured across standard metrics on various substructure prediction tasks. Indeed, we find that for predicting the presence of both 86 molecular fragments arising purely from structural information, as well as 55 structural alerts associated with molecule toxicity, LaPool globally outperforms other pooling models and GIN for the F1 and ROC-AUC metrics (micro-averaged to deal with high class imbalance). In particular, on the harder and extremely imbalanced molecular alerts prediction task, all models performed poorly compared to LaPool, suggesting that the hierarchical representation learned by LaPool helps to achieve a better understanding of the molecular graphs.

heading: 4.2 EXPERIMENTS ON STANDARD GRAPH CLASSIFICATION BENCHMARKS
text: In addition to evaluating molecular structural understanding of the pooling models, we benchmark our model on molecular toxicity prediction using the TOX21 dataset (Council et al., 2007). We further conduct experiments on non-molecular benchmark graph datasets (DD, PROTEINS, FRANKENSTEIN), which usually contain larger and often denser graphs compared to molecular graphs (see Supplemental section C for dataset statistics). For TOX21, we report the test ROC-AUC averaged over 5-folds (following the 80-10-10 split proportion used in Wu et al. (2018)), while we follow prior work (Ying et al., 2018) on the remaining datasets by reporting the best accuracy on a 10-fold cross-validation.
As shown in Table 2, LaPool outperforms all other approaches on the well known TOX21 dataset and on the PROTEINS and FRANKENSTEIN, both of which contain non-molecular graphs with size similar to the TOX21 molecules. In particular, on the PROTEINS dataset LaPool achieved an
accuracy of 83.83, representing a significant gap relative to DiffPool, its closest competitor, at 77.25. This suggests that the LaPool method is not restricted to molecular data but has broad applicability, especially in the context of sparse graph classification.

heading: 4.3 MOLECULAR GENERATION
text: We aim to showcase LaPool’s utility in drug discovery by demonstrating that it can be leveraged to generate molecules. In previous work, GANs and VAEs were used to generate either string or graph representations of molecules. Here, we use the GAN-based Wasserstein Auto-Encoder recently proposed in (Tolstikhin et al., 2017) to model the data distribution of molecules (see Figure 1 in Supplemental Material). For the encoder, we use a similar network architecture as in our supervised experiments. The decoder and discriminator are simple MLPs, with complete architectural details provided in Supplemental Section A.4. Although the encoder is permutation invariant, the decoding process may not be. To force the decoder to learn a single graph ordering, we use a canonicalization algorithm (Schneider et al., 2015) that reorders atoms to ensure a unique graph for each molecule. We further improve the robustness of our generative model to node permutations by computing the reconstruction loss using a permutation-invariant embedding, parameterized by a GIN, on both the input and reconstructed graphs (see Supplemental Section A.4.2). We find that such a formulation improves the reconstruction loss and increases the ratio of valid molecules generated.
Dataset and Baseline Models Following previous work on molecular generation, we evaluate our generative model with an encoder enhanced by the LaPool layer (referred to as WAE-LaP) on the QM9 molecular dataset (Ramakrishnan et al., 2014). This dataset contains 133,885 small drug-like organic compounds with up to 9 heavy atoms (C, O, N, F). We compare WAE-LaP to alternatives within our WAE framework where either no pooling is used (WAE-GNN) or where DiffPool is used as the pooling layer (WAE-Diff). Our results are also compared to previous results on the same dataset, including Grammar-VAE, GraphVAE, and MolGAN.
Evaluation Metrics We measure the performance of the generative model using metrics standard in the field: validity (proportion of valid molecules from generated samples), uniqueness (proportion of unique molecules generated), and novelty (proportion of generated samples not found in the training set). All metrics were computed on a set of 10,000 generated molecules.
As shown in Table 3, WAE-LaP generated the most valid and unique molecules compared to all WAEbased generative models with slightly lower but similar novelty. Although MolGAN performed best on the novelty metric, it has among the lowest percentage of unique molecules. We hypothesize that the decrease in novelty observed with LaPool might be a result of its pooling mechanism encouraging fragment novelty during sampling, thus limiting novelty resulting from rearrangement at the atom level. Nevertheless, as all WAE-based methods produced similar proportions of novel molecules, our results still suggest that combining LaPool with other generative approaches could improve the uniqueness and validity of generated compounds. We therefore conclude that the pooling performed by LaPool can improve molecular graph representation, which is crucial in a generative setting.

heading: 4.4 IMPROVED INTERPRETABILITY
text: To better understand the insights provided by LaPool, we conduct model interpretability experiments on molecular fragment prediction. Here we refer to interpretability as the degree to which a human (in this context, a medicinal chemist) can understand the cause of the model’s decision (Miller, 2018). This explains our focus on fragment prediction since in that setting an “interpretable model” that achieves high performance would need to first understand the graph structure and the relationship between nodes.

heading: Graph U-net
text: We first investigate the behavior of LaPool and DiffPool by analyzing the clustering made at the pooling layer level. This comparison was limited to DiffPool since Graph U-net does not perform a node clustering. We argue that an interpretable pooling layer should preserve the overall graph structure after pooling and produce meaningful clusters that could provide insight into the contribution of each molecular subgraph from the perspective of an expert chemist. While defining what is meaningful is inherently subjective, we attempt to shed light on these models by observing their behavior in the drug discovery domain, using our understanding of chemical structure as reference.
We show in Figure 2 that LaPool is able to coarsen the molecular graphs into sparsely connected graphs, which can be interpreted as the skeleton of the molecules. Indeed, the data-driven dynamic segmentation it performed is akin to chemical fragmentation (Gordon et al., 2011). In contrast, DiffPool’s cluster assignment is more uniform across the graph, leading to densely connected coarsened graphs which are less interpretable from a chemical viewpoint. In particular, it fails in the presence of molecular symmetry, as it encourages the mapping of nodes with similar features to the same clusters. This is illustrated in both example (c) which shows how DiffPool creates a fully connected graph from an originally disconnected graph, and example (b) which shows how symmetric elements, despite being far from each other, are assigned identically. On the other hand, we observe that Graph U-net ignores the graph structure, typically disconnecting it. It also appears very biased toward selecting atoms in similar environment to ones already selected. Such failures are not present when using LaPool, since the dynamic centroid selection and the subsequent distance regularization enforce preservation of the molecular graph structure. A typical failure case for LaPool is seen in (e) and corresponds to a missing centroid node in a given region of the graph, which results in a soft assignment of the region to multiple clusters. However, this behavior is inherent to most DiffPool samples since the fixed number of clusters and the inability to consider node distance cannot account for the diversity of molecular datasets.

heading: C25
text: In addition to assessing the quality of clustering performed by LaPool and DiffPool, we attempt to directly target interpretability by computing an explanation of each model decision and comparing it to a ground truth. We design a simple experiment in which we predict the presence of either Epoxide, Thioepoxide, or Aziridine substructures (denoted by the molecular pattern “C1[O,S,N]C1”), that are indicative of molecular toxicity. Interpretability is therefore defined as the accuracy of the importance attributed by each model to relevant substructures of the input molecules, given the presence of the underlying ground truth fragment we wish to predict. Similar to (Pope et al., 2019), we adapt an existing explainability method for CNNs to GNNs. Specifically, we choose to compute the integrated gradient (Sundararajan et al., 2017) over the input node features due to its stability and robustness in the presence of zero-value features (see Supplemental section C for discussion and alternate approach). Next, we derive an importance score for each node using the L1-norm of the feature attribution map for the node. By both qualitatively observing samples from the data (Figure 3), and by measuring the PR-AUC over the computed importance values given the ground truth to assess the ability to distinguish between important and non-important nodes (Table 4), we find LaPool to more robustly identify the salient structure, resulting in improved overall interpretability. An interesting outcome of this experiment is the performance of Graph U-net, ranked second best. Such performance is a direct result of using a cluster size large enough to cover the toxic fragment size.

heading: 5 CONCLUSION
text: In this work, we have proposed LaPool, a differentiable and robust pooling operator for molecular and sparse graphs that considers both node information and graph structure. In doing so, we have proposed a method which is able to identify important substructures of a graph by leveraging the graph Laplacian. In contrast with previous work, this method retains the connectivity structure and
feature information of the graph during the coarsening procedure, while encouraging nodes belonging to the same substructure to be mapped together in the coarsened graph.
Incorporating the proposed pooling layer into existing graph neural networks, we have demonstrated that the enforced hierarchization allows for the capture of a richer and more relevant set of features at the graph-level representation. We discussed the performance of LaPool relative to existing graph pooling layers and demonstrated on both molecular graph classification and generation benchmarks that LaPool outperforms existing graph pooling modules and produces more interpretable results. In particular, we argue that the molecular graph segmentation performed by LaPool provides greater insight into molecular activity and that the associated properties can be leveraged in drug discovery. Finally, we show that although LaPool was designed for molecular graphs, it generalizes well to other graph types. In future work, we aim to investigate how additional sources of information such as edge features could be incorporated into the graph pooling process.

heading: A NETWORK ARCHITECTURE AND TRAINING PROCEDURE
text: Below, we describe the network architecture and the training procedure used for both supervised and generative experiments.

heading: A.1 EDGE ATTRIBUTES
text: Part of the work presented assumes the absence of edge attributes in the graphs. However, in molecular graphs, the nature of a bond between two atoms plays an important role regarding activity and property. As such, edge types should be considered, especially in generative models. To consider this, we add to our network an initial Edge-GC layer described in the following.
Let G = 〈V,E,X〉 be an undirected molecular graph, such that E = [E1, . . . Ek] ∈ {0, 1}e×n×n where n is the number of nodes in the graph and e is the number of possible edge. We have that∑
1≤i≤e
E::i = A ∈ {0, 1}n×n (7)
where A is the adjacency matrix of the graph. The Edge GC layer is simply defined as :
Y = MΘ1(E1, X)‖ . . . ‖MΘe(Ee, X) (8)
where ‖ is the concatenation operator on the node feature dimension and MΘ1 are graph neural networks parameterized to learn different features for each edge type. A new graph defined as G′ = 〈V,A, Y 〉 can then be feed into the subsequent layers of the network.

heading: A.2 ATOM AND EDGE FEATURES
text: In our experiments, the initial node feature tensor is represented by a one-hot encoding of atoms (ignoring hydrogens) within the respective datasets and additional properties such as the atom implicit valence, its formal charge, number of radical electrons and whether it is in a molecular ring. For edge attributes, we consider the single, double and triple bond, which were enough to cover all molecules in the datasets, given that feature extraction was preceded by kekulization of molecules.

heading: A.3 SUPERVISED EXPERIMENTS
text: In all of our supervised experiments, we use a graph convolution module consisting of two graph convolutional layers of 128 channels each with ReLU activation; followed by an optional hierarchical graph pooling layer; then two additional graph convolution layers (64) with skip connection to introduce jumping knowledge and a gated global graph pooling layer (Li et al., 2015) to yield a graph-level representation. This is further followed by one fully connected layers (128) with batch normalization and ReLU activation, finalized by a linear output layer with appropriate activation for the task readouts. Notice that we used one pooling layer, since no noticeable improvement was observed when using more in our experimental setting.
For DiffPool, we performed a hyperparameter search to find the optimal number of clusters (12.5%, 25%, 50% of the maximum number of nodes in the batch (Ying et al., 2018)). A similar search is also performed for the Graph U-net pooling layer. For LaPool, we consider the same number of clusters and the dynamic node seelction. We also performed a grid search over the window size k used as regularization to prevent nodes from mapping to centroids that are more than k-hop away as an alternative to the distance-regularized version. The grid search was performed for k ∈ {1, 2, 3}. For the supervised experiments, we use a batch size of 64 and train the networks for 100 epochs, with early stopping.

heading: A.4 GENERATIVE MODELS

heading: A.4.1 WAE MODEL
text: We use a Wasserstein Auto-Encoder (WAE) as our generative model (see Figure S1. The WAE minimizes a penalized form of the Wasserstein distance between a model distribution and a target
G NN
L ay
er
EG
NN L
ay er
G ra
ph P
oo lin
g
G -S
um Po
ol
FC L
ay er
O O
NH
SE
G O
O
NH
0/1 D
G NN
L ay
er
... ...
l1
l2
a) Adversarial autoencoder (AAE) b) Encoder network
Figure S1: Model architecture for the generative model. (a) We use a WAE, in which a generator (auto-encoder) progressively learns the true molecular data distribution. (b) Architecture used for the encoder network.
distribution, and has been shown to improve learning stability. As described in (Tolstikhin et al., 2017), we aim to minimize the following objective:
inf Q(Z|X)∈Q EPXEQ(Z|X)[cost(X,G(Z)] + λDZ(QZ , PZ) (9)
where Q is any nonparametric set of probabilistic encoders, DZ is the Jensen-Shannon divergence between the learned latent distribution QZ and prior PZ , and λ > 0 is a hyperparameter. DZ is estimated using adversarial training (discriminator).
For our generative model, the encoder follows a similar structure as the network used for our supervised experiments, with the exception being that the network now learns a continuous latent space qΨ(z|G) given a set of input molecular graphs G = {G1, · · · , Gn}. More precisely, it consists of one edge graph layer, followed by two GCs (32 channels each), an optional hierarchical graph pooling, two additional GC layers (64, 64), then one global sum pooling step (128) and two fully connected layers (128), meaning the molecular graphs are embedded into a latent space of dimension 128. Instead of modeling the node/edge decoding with an autoregressive framework as done in recent works (You et al., 2018b; Assouel et al., 2018; Li et al., 2018d) to capture the interdependency between them, we used a simple MLP that takes the latent code z as input an pass it through two fully connected layers (128, 64). The output of those layers is used as shared embedding for two networks: one predicting the upper triangular entries of the edge tensor, and the second predicting the node features tensor. This results in faster convergence.
For the discriminator, we use a simple MLP that predicts whether the latent code comes from a normal prior distribution z N (0, 1). This MLP is constituted by two stacked FCLs (64, 32) followed by an output layer with sigmoid activation. As in (Kadurin et al., 2017), we do not use batch-normalization, since it resulted in a mismatch between the discriminator and the generator.
All models use the same basic generative architecture, with the only difference being the presence of a pooling-layer and its associated parameters. For DiffPool, we fixed the number of cluster to three, while for LaPool, we use the distance-based regularization.

heading: A.4.2 RECONSTRUCTION LOSS
text: For each input molecular graph G = 〈V,E,X〉, the decoder reconstruct a graph G̃ = 〈Ṽ , Ẽ, X̃〉. Since we use a canonical ordering (available in RDKit) to construct G from the SMILES representation of molecules, the decoder is forced to learn how to generate a graph under this order. Therefore, the decoding process is not necessarily able to consider permutations on the vertices set, and generation of isomorphic graphs will be heavily penalized in the reconstruction loss. In (Simonovsky and Komodakis), the authors use an expensive graph matching procedure to overcome that limitation. We argue that it suffices to compute the reconstruction loss on γ(G) and γ(G̃), where γ is a permutation invariant embedding function. As a heuristic, we used a Graph Isomorphism Network (GIN), with weights fixed to 1, in order to approximate the Weisfeiler-Lehman graph isomorphism test (see (Xu
et al., 2018) for more details). In particular, we use an edge-aware GIN layer (see section A.1) to embed both G and G̃. The reconstruction loss is then defined as:
Lrec = 1 |V | ∑ i (γ(G)i − γ(G̃)i)2 (10)
Our experiments show that this loss function was able to produce a higher number of valid molecules, although we speculate that such a heuristic might prove harder to optimize on datasets with larger graphs.

heading: A.4.3 TRAINING PROCEDURE
text: The QM9 dataset was split into a train (60%), valid (20%) and a hold-out test dataset (20%). Note that only 25% of the training set is sampled during each epoch (batch size 32). The generator network (encoder-decoder) and the discriminator network are trained independently, using the Adam optimizer Kingma and Ba (2014) with an initial learning rate of 1e− 4 for the generator and 1e− 3 for the discriminator. During training, we slowly reduce the learning rate by a factor of 0.5, for the generator, on plateau. To stabilize the learning process and prevent the discriminator from becoming "too good" at distinguishing the true data distribution from the prior, we train the generator two times more often.

heading: B GENERATED MOLECULES
text: Below, we highlight a few molecules generated by WAE-LaP on the QM9 dataset.

heading: OH NH2

heading: C DATASET STATISTICS
text: TOX21 DD PROTEINS FRANKENSTEIN
Avg. nodes 18.51 284.32 39.05 16.83 Avg. edges 19.23 715.66 72.82 17.88 #Graphs 8014 1178 1113 4337 #Classes 12 2 2 2

heading: D NODE IMPORTANCE INTERPRETABILITY SCORE
text: In addition to assessing the quality of clustering performed by LaPool and DiffPool, we attempt to measure the interpretability of their predictions. We consider a setting in which the goal is to predict the presence of either Epoxides, Thioepoxides or Aziridines substructures in molecular graphs. These
three fragments correspond to structural alerts that are often indicative of molecular toxicity. We attempt to identify the relational inductive bias used by each model during prediction. In our setting, we define the interpretability of a model as its ability to focus on nodes that are directly relevant to the structural alerts and leverage that information for its prediction. In other words, we expect the most important nodes for the model prediction to correspond to nodes that are part of the structural alerts. We measure the importance of each atom toward the model prediction using the Integrated gradient method. Briefly, we compute perturbation of node and edge attributes over a continuous spectrum, then integrate the gradient of each of the model loss with respect to both the perturbed adjacency matrix and node features. Similar to saliency maps, we then take the sum of the absolute integrated gradients over each node as an approximate attribution score for the nodes. Finally, we compute the interpretability score using the Precision-Recall AUC between measured importance and ground truth which is defined by a binary mask of nodes that are part of the structural alerts. The PR-AUC allows us to assess the node importance separation capacity of each model while taking imbalance into account. We only focus on positive predictions for each model. As an alternative to the Integrated Gradient, we also measure the interpretability score using Guided BackPropagation (see Table 6)

heading: E SIGNAL PRESERVATION THROUGH LAPLACIAN MAXIMA
text: We illustrate here on a 1-d signal S, how using the Laplacian maxima serves to retain the most prominent regions of the graph signal, after smoothing (Figure S3). We measure the energy conservation after downsampling: δE(S) = E(S)− E(Sdown) of the 1-d signal energy to highlight why selecting the Laplacian maxima allow reconstructing the signal with a low error when compared to the minimum Laplacian (which focuses on low frequencies). The energy ES of a discrete signal yi is defined in (11), and is similar to the energy of a wave in a physical system (without the constants).
ES = ∑ i |yi|2 (11)
0 5 10 15 20 25 Node number
0
2
4
Si gn
al in
te ns ity Random smooth 1D signal - 25 nodes, ES = 71.1Maximum Laplacian pooling - 8 leaders, ES = 72.5 Minimum Laplacian pooling - 7 leaders, ES = 76.3
Figure S3: Comparison of maximum/minimum Laplacian pooling for a random and smoothed signal on a 1D graph with 25 nodes. The graph energy ES is indicated.
To mimic the molecular graph signal at the pooling stage, the given signal is built from an 8-terms random Fourier series with added Gaussian noise, then smoothed with 2 consecutive neighbor average smoothing. For the pooling methods, a linear interpolation is used to cover the same signal space before computing ES . As expected, the maxima Laplacian selection requires a fewer number of samples for signal reconstruction and energy preservaton. It also significantly outperforms minima selection.

Review summaries:
The authors propose a new pooling layer, LaPool, for hierarchical graph representation learning (Ying et al., 2019) by clustering nodes around centroids that are selected based on "signal intensity variation". The signal intensity variation of node x is defined as sum_{y in HOP(x, h)} ||x - y||  where HOP(x, h) is the set of nodes reachable from x within h hops. Once top k maximizers are selected as centroids (k can be predetermined or dynamically chosen), a sparse cluster assignment distribution is computed for each node using sparsemax (Laha et al., 2018), and the affinity matrix and the node embeddings are coarsened as in Ying et al. (2019). The authors show that LaPool can improve performance in various graph-related tasks over baselines and generate interpretable clusters. 

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.