Original: T OWARDS I NTERPRETABLE M OLECULAR G RAPH R EPRESENTATION L EARNING
Revised: T OWARDS I NTERPRETABLE M OLECULAR G RAPH R EPRESENTATION L EARNING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Recent work in graph neural networks (GNNs) has led to improvements in molecular activity and property prediction tasks.
Revised: Recent work in graph neural networks (GNNs) has led to improvements in molecular activity and property prediction tasks.

Original: Unfortunately, GNNs often fail to capture the relative importance of interactions between molecular substructures, in part due to the absence of efﬁcient intermediate pooling steps.
Revised: Unfortunately, GNNs often fail to capture the relative importance of interactions between molecular substructures, in part due to the absence of efﬁcient intermediate pooling steps.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We then demonstrate the improved interpretability achieved with LaPool using both qualitative and quantitative assessments, highlighting its potential applications in drug discovery.
Revised: We then demonstrate the improved interpretability achieved with LaPool using both qualitative and quantitative assessments, highlighting its potential applications in drug discovery.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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).
Revised: 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).

Original: To this end, graph neural networks (GNN) falling into either spectral-based or spatial-based approaches have been proposed.
Revised: To this end, graph neural networks (GNN) falling into either spectral-based or spatial-based approaches have been proposed.

Original: Spectral methods deﬁne the graph convolution (GC) as a ﬁltering operator of the graph signal (Defferrard et al., 2016), while spatial methods deﬁne the GC as a message passing and aggregation across nodes (Henaff et al., 2015;
Revised: Spectral methods deﬁne the graph convolution (GC) as a ﬁltering operator of the graph signal (Defferrard et al., 2016), while spatial methods deﬁne the GC as a message passing and aggregation across nodes (Henaff et al., 2015;

Original: Xu et al., 2018; Jin et al., 2018).
Revised: Xu et al., 2018; Jin et al., 2018).

Original: In drug discovery, GNNs have been very successful across several molecular graph classiﬁcation and generation tasks.
Revised: In drug discovery, GNNs have been very successful across several molecular graph classiﬁcation and generation tasks.

Original: In particular, they outperform predetermined molecular ﬁngerprints 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).
Revised: In particular, they outperform predetermined molecular ﬁngerprints 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).

Original: However, the node feature update performed by GNNs introduces important limitations.
Revised: However, the node feature update performed by GNNs introduces important limitations.

Original: 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).
Revised: 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).

Original: This limits the network’s depth and restricts the receptive ﬁeld of the vertices in the graph to a few-hop neighborhood, which is insufﬁcient to properly capture local structures, relationships between nodes, and subgraph importance in sparse graphs such as molecules.
Revised: This limits the network’s depth and restricts the receptive ﬁeld of the vertices in the graph to a few-hop neighborhood, which is insufﬁcient to properly capture local structures, relationships between nodes, and subgraph importance in sparse graphs such as molecules.

Original: For example, at least three consecutive GC layers are needed for atoms at the opposite side of a benzene ring to exchange information.
Revised: For example, at least three consecutive GC layers are needed for atoms at the opposite side of a benzene ring to exchange information.

Original: 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.
Revised: 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.

Original: To cope with these limitations, graph coarsening (pooling) methods have been proposed to reduce graph size and enable long-distance interaction between nodes.
Revised: To cope with these limitations, graph coarsening (pooling) methods have been proposed to reduce graph size and enable long-distance interaction between nodes.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: The primary contributions of this paper are summarized below:
Revised: The primary contributions of this paper are summarized below:

Original: • We propose a novel and differentiable pooling module (LaPool) that can be incorporated into existing GNNs to yield more expressive networks for molecular data.
Revised: • We propose a novel and differentiable pooling module (LaPool) that can be incorporated into existing GNNs to yield more expressive networks for molecular data.

Original: • We propose a graph structure understanding dataset for benchmarking GNNs that is based on molecular substructure prediction.
Revised: • We propose a graph structure understanding dataset for benchmarking GNNs that is based on molecular substructure prediction.

Original: • We show that LaPool outperforms recently proposed graph pooling layers on both discriminative and generative molecular graph benchmarks, while also remaining competitive on other graph benchmarks.
Revised: • We show that LaPool outperforms recently proposed graph pooling layers on both discriminative and generative molecular graph benchmarks, while also remaining competitive on other graph benchmarks.

Original: • We highlight the improved interpretability achieved by LaPool using both qualitative and quantitative assessments.
Revised: • We highlight the improved interpretability achieved by LaPool using both qualitative and quantitative assessments.

Original: 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.
Revised: 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.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: In this section, we ﬁrst introduce related work on graph pooling, then provide an overview of techniques used in computational drug discovery to put our work into context.
Revised: In this section, we ﬁrst introduce related work on graph pooling, then provide an overview of techniques used in computational drug discovery to put our work into context.

Original: As our focus herein is on graph pooling, we refer readers to (Wu et al., 2019) for an overview of recent progress in GNNs.
Revised: As our focus herein is on graph pooling, we refer readers to (Wu et al., 2019) for an overview of recent progress in GNNs.

Original: In traditional GNN architectures, global sum/average/max-pooling layers have been used to aggregate node embeddings into a graph-level representation.
Revised: In traditional GNN architectures, global sum/average/max-pooling layers have been used to aggregate node embeddings into a graph-level representation.

Original: Recently, more sophisticated methods have been proposed.
Revised: Recently, more sophisticated methods have been proposed.

Original: For example, Li et al.
Revised: For example, Li et al.

Original: (2015) use a gated mechanism, Zhang et al.
Revised: (2015) use a gated mechanism, Zhang et al.

Original: (2018) proposed
Revised: (2018) proposed

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Consequently, hierarchical graph pooling methods have been proposed.
Revised: Consequently, hierarchical graph pooling methods have been proposed.

Original: They act by reducing graph size and increasing node receptive ﬁelds without increasing network depth.
Revised: They act by reducing graph size and increasing node receptive ﬁelds without increasing network depth.

Original: However, in contrast to the regular structure of images, graphs are irregular and complex, making it challenging to properly pool nodes together.
Revised: However, in contrast to the regular structure of images, graphs are irregular and complex, making it challenging to properly pool nodes together.

Original: 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).
Revised: 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).

Original: More recently, differentiable hierarchical graph pooling layers have been proposed.
Revised: More recently, differentiable hierarchical graph pooling layers have been proposed.

Original: Ying et al.
Revised: Ying et al.

Original: (2018) proposed DiffPool, a pooling layer that performs a similarity-based node clustering using a soft afﬁnity matrix learned by a GNN.
Revised: (2018) proposed DiffPool, a pooling layer that performs a similarity-based node clustering using a soft afﬁnity matrix learned by a GNN.

Original: 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.
Revised: 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.

Original: In computer-aided drug discovery, methods such as virtual screening and de novo drug design serve as efﬁcient complements to physical high-throughput screening of large molecular libraries.
Revised: In computer-aided drug discovery, methods such as virtual screening and de novo drug design serve as efﬁcient complements to physical high-throughput screening of large molecular libraries.

Original: 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).
Revised: 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).

Original: Importantly, data-driven virtual screening approaches that leverage advances in deep learning, rather than pre-determined features such as molecular ﬁngerprints (Rogers and Hahn, 2010) and SMILES string representations, have been shown to dramatically improve prediction accuracy (Kearnes et al., 2016; Wu et al., 2018).
Revised: Importantly, data-driven virtual screening approaches that leverage advances in deep learning, rather than pre-determined features such as molecular ﬁngerprints (Rogers and Hahn, 2010) and SMILES string representations, have been shown to dramatically improve prediction accuracy (Kearnes et al., 2016; Wu et al., 2018).

Original: 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.
Revised: 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.

Original: The ﬁrst 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.
Revised: The ﬁrst 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.

Original: Graph generative models have since been developed (e.g.
Revised: Graph generative models have since been developed (e.g.

Original: JT-VAE (Jin et al., 2018), GraphVAE (Simonovsky and Komodakis), MolGAN (De Cao and Kipf, 2018), MolMP (Li et al., 2018b), etc.)
Revised: JT-VAE (Jin et al., 2018), GraphVAE (Simonovsky and Komodakis), MolGAN (De Cao and Kipf, 2018), MolMP (Li et al., 2018b), etc.)

Original: and have been shown to improve the validity and novelty of generated molecules.
Revised: and have been shown to improve the validity and novelty of generated molecules.

Original: In addition, these methods allow conditional molecule generation via Bayesian optimization or reinforcement learning (Jin et al., 2018;
Revised: In addition, these methods allow conditional molecule generation via Bayesian optimization or reinforcement learning (Jin et al., 2018;

Original: Olivecrona et al., 2017; Assouel et al., 2018;
Revised: Olivecrona et al., 2017; Assouel et al., 2018;

Original: Li et al., 2018d; You et al., 2018a).
Revised: Li et al., 2018d; You et al., 2018a).

Original: In this work, we are mainly interested in the impact of molecular representation on generative performance as opposed to the optimization procedure itself.
Revised: In this work, we are mainly interested in the impact of molecular representation on generative performance as opposed to the optimization procedure itself.

Original: G RAPH L APLACIAN P OOLING
Revised: G RAPH L APLACIAN P OOLING

Original: O O
Revised: O O

Original: S
Revised: S

Original: N
Revised: N

Original: C H C H
Revised: C H C H

Original: C H
Revised: C H

Original: C
Revised: C

Original: C H
Revised: C H

Original: C O
Revised: C O

Original: O
Revised: O

Original: S
Revised: S

Original: N
Revised: N

Original: C H C H
Revised: C H C H

Original: C H
Revised: C H

Original: C
Revised: C

Original: C H
Revised: C H

Original: C
Revised: C

Original: O O
Revised: O O

Original: S
Revised: S

Original: N
Revised: N

Original: C H C H
Revised: C H C H

Original: C H
Revised: C H

Original: C
Revised: C

Original: C H
Revised: C H

Original: C
Revised: C

Original: [Figure]
Revised: [Figure]

Original: A reliable pooling operator should maintain the overall structure and connectivity of a graph.
Revised: A reliable pooling operator should maintain the overall structure and connectivity of a graph.

Original: LaPool achieves this by taking into account the local structure deﬁned by the neighborhood of each node.
Revised: LaPool achieves this by taking into account the local structure deﬁned by the neighborhood of each node.

Original: As shown in Figure 1, the method uses a standard GC layer with a centroid selection and a follower selection step.
Revised: As shown in Figure 1, the method uses a standard GC layer with a centroid selection and a follower selection step.

Original: First, the centroids of the graph are selected based on the local signal variation (see Section 3.2).
Revised: First, the centroids of the graph are selected based on the local signal variation (see Section 3.2).

Original: Next, LaPool learns an afﬁnity matrix C using a distance normalized attention mechanism to assign all nodes of the graph to the centroids (see Section 3.3).
Revised: Next, LaPool learns an afﬁnity matrix C using a distance normalized attention mechanism to assign all nodes of the graph to the centroids (see Section 3.3).

Original: Finally, the afﬁnity matrix is used to coarsen the graph.
Revised: Finally, the afﬁnity matrix is used to coarsen the graph.

Original: These steps are detailed below.
Revised: These steps are detailed below.

Original: 3. P RELIMINARIES
Revised: 3. P RELIMINARIES

Original: Notation Let G =
Revised: Notation Let G =

Original: (cid:104) V, A, X (cid:105) be an undirected graph, where V = (cid:8) v 1 , . . .
Revised: (cid:104) V, A, X (cid:105) be an undirected graph, where V = (cid:8) v 1 , . . .

Original: v n (cid:9) is its vertex set,
Revised: v n (cid:9) is its vertex set,

Original: A ∈ { 0 , 1 } n × n denotes its adjacency matrix, and X = (cid:2) x 1 , . . .
Revised: A ∈ { 0 , 1 } n × n denotes its adjacency matrix, and X = (cid:2) x 1 , . . .

Original: x n (cid:3) (cid:124) ∈ R n × d is the node feature matrix with each node v i having d -dimensional feature x i .
Revised: x n (cid:3) (cid:124) ∈ R n × d is the node feature matrix with each node v i having d -dimensional feature x i .

Original: X can also be viewed as a d -dimensional signal on G (Shuman et al., 2012).
Revised: X can also be viewed as a d -dimensional signal on G (Shuman et al., 2012).

Original: Without loss of generality we may assume a ﬁxed ordering of the nodes that is respected in V , A , and X .
Revised: Without loss of generality we may assume a ﬁxed ordering of the nodes that is respected in V , A , and X .

Revised: The neighborhood of radius h (or h -hop) neighborhood of a node v i ∈ V is the set of nodes separated from v i by a path of length at most h and is denoted by

Original:  Graph Signal For any graph G , itsnormalized graph Laplacian matrix L is deﬁned as L =I − D −1 A , where D is a diagonal matrix with D ii being the degree of node v i .
Revised: N h ( v i ) . For simplicity, we will use N ( v i ) to refer to the set of nodes adjacent to v i . Graph Signal For any graph G , its graph Laplacian matrix L is deﬁned as L = D − A , where D is a diagonal matrix with D i,i being the degree of node v i .

Original: The graph Laplacian is a difference operator and can be used to deﬁne the smoothness s ( X ) (the extent at which the signal changes between connected nodes) of a signal X on G .
Revised: The graph Laplacian is a difference operator and can be used to deﬁne the smoothness s ( X ) (the extent at which the signal changes between connected nodes) of a signal X on G .

Original: For 1-dimensional signal f = [ f 1 , . . .
Revised: For 1-dimensional signal f = [ f 1 , . . .

Original: , f n ] :
Revised: , f n ] :

Original: [Equation]
Revised: [Equation]

Original: Graph Neural Networks We consider GNNs that act in the graph spatial domain as message passing (Gilmer et al., 2017).
Revised: Graph Neural Networks We consider GNNs that act in the graph spatial domain as message passing (Gilmer et al., 2017).

Original: We focus on the Graph Isomorphism Network (GIN)
Revised: We focus on the Graph Isomorphism Network (GIN)

Original: (Xu et al., 2018), which uses a SUM-aggregator on messages received by each node to achieve a better understanding of the graph structure:
Revised: (Xu et al., 2018), which uses a SUM-aggregator on messages received by each node to achieve a better understanding of the graph structure:

Original: [Equation]
Revised: [Equation]

Original: where M l Θ is a neural network with trainable parameters
Revised: where M l Θ is a neural network with trainable parameters

Original: Θ , x i is the feature vector for node v i , v j ∈ N ( v i ) are the neighbours of v i and l is the layer number.
Revised: Θ , x i is the feature vector for node v i , v j ∈ N ( v i ) are the neighbours of v i and l is the layer number.

Original: Notice the term
Revised: Notice the term

Original: A
Revised: A

Original: i,j that takes into account the edge weight between nodes v i and v j when A is not binary.
Revised: i,j that takes into account the edge weight between nodes v i and v j when A is not binary.

Original: In this work, we focus mainly on molecular graphs in supervised settings where, given a molecule m and its corresponding molecular graph G m , we aim to predict some properties of m .
Revised: In this work, we focus mainly on molecular graphs in supervised settings where, given a molecule m and its corresponding molecular graph G m , we aim to predict some properties of m .

Original: 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.
Revised: 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.

Original: 3. G RAPH D OWNSAMPLINGVIA B AND -P ASS F ILTERING
Revised: 3. G RAPH D OWNSAMPLING

Original: This section details how LaPool downsamples the original graph by selecting a set V C of nodes as centroids after l consecutive GC layers.
Revised: This section details how LaPool downsamples the original graph by selecting a set V C of nodes as centroids after l consecutive GC layers.

Original: Centroid Selection For any given vertex v i , we can deﬁne a local measure s i ofsignal intensity  variation around v i .
Revised: Centroid Selection For any given vertex v i , we can deﬁne a local measure s i of intensity of signal variation around v i .

Original: As s i measures how different the signal residing at a node v
Revised: As s i measures how different the signal residing at a node v

Original: i is fromthe average of its neighbours reachable within a h − hop, we are interested in the set V C of nodes that have a signal intensity  variation s i greater than  the average in their neighborhood, which can be obtained from the Laplacian.
Revised: i is from its neighbours, we are interested in the set V C of nodes that have an intensity of signal variation s i greater than their neighborhood. In this work, we use a deﬁnition 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:

Original: [Equation]
Revised: [Equation]

Revised: where top V ( L · S, k ) corresponds to the top k nodes with the greatest intensity of signal variation, and where || · || 2 , R d corresponds to taking the vector norms over the d-dimensional rows of LX .

Revised: Instead of using the direct neighbours, one can also generalize the computation of S in equation 3 to an h − hop neighbourhood by taking powers of the Laplacian.

Original: Observe that the GC layers preceding each pooling step perform a smoothing of the graph signal and thus act as a low-pass ﬁlter.
Revised: Observe that the GC layers preceding each pooling step perform a smoothing of the graph signal and thus act as a low-pass ﬁlter.

Original: Combined with the high-pass ﬁlter of Eq.
Revised: Eq.

Original: (3), it results in aband-pass ﬁltering of X that attenuates low and high-frequency noise, yet retains the important signal information.
Revised: (3) emphasizes instead the high variation regions, resulting overall in a ﬁltering of X that attenuates low and high-frequency noise, yet retains the important signal information.

Original: 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.
Revised: 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.

Original: 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., 2015b;a).
Revised: 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).

Original: For a 1D signal, LaPool selects points, usually near the maxima/minima, where the derivative changes the most and is hardest to interpolate linearly.
Revised: For a 1D signal, LaPool selects points, usually near the maxima/minima, where the derivative changes the most and is hardest to interpolate linearly.

Original: For molecular graphs, this often corresponds to sampling a subset of nodes critical for reconstructing the original molecule.
Revised: For molecular graphs, this often corresponds to sampling a subset of nodes critical for reconstructing the original molecule.

Original: Dynamic Selection of the Centroids The method presented in Eq.
Revised: Dynamic Selection of the Centroids The method presented in Eq.

Original: (3) implies the selection of k centroids.
Revised: (3) implies the selection of k centroids.

Original: 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
Revised: 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

Original: s i greater than its neighbours s 
Revised: s i greater than its neighbours s j :

Original: [Equation]
Revised: [Equation]

Original: 3. L EARNING THE N ODE - TO -C LUSTER A SSIGNMENT M ATRIX
Revised: 3. L EARNING THE N ODE - TO -C LUSTER A SSIGNMENT M ATRIX

Original: Once the set V C of centroid nodes is determined, we compute a mapping of the remaining “follower” nodes V F = V \ V C into the new clusters formed by the nodes in V C .
Revised: Once the set V C of centroid nodes is determined, we compute a mapping of the remaining “follower” nodes V F = V \ V C into the new clusters formed by the nodes in V C .

Original: This mapping gives the cluster assignment C = (cid:2) c 1 , ...
Revised: This mapping gives the cluster assignment C = (cid:2) c 1 , ...

Original: c n (cid:3) (cid:124) ∈ [0 , 1] n × m
Revised: c n (cid:3) (cid:124) ∈ [0 , 1] n × m

Original: s.t. ∀
Revised: s.t. ∀

Original: i : 1c (cid:124) i = 1 , where each row c i corresponds to the afﬁnity of node v i towards each of the m clusters in V C .
Revised: i : 1c (cid:124) i = 1 , where each row c i corresponds to the afﬁnity of node v i towards each of the m clusters in V C .

Original: Let X ( l ) be the node embedding matrix at an arbitrary layer and X ( l ) C the embedding of the “centroids”.
Revised: Let X ( l ) be the node embedding matrix at an arbitrary layer and X ( l ) C the embedding of the “centroids”.

Original: We compute C using a soft-attention mechanism (Graves et al., 2014) measured by the cosine similarity between X ( l ) and X ( l ) C :
Revised: We compute C using a soft-attention mechanism (Graves et al., 2014) measured by the cosine similarity between X ( l ) and X ( l ) C :

Original: [Equation]
Revised: [Equation]

Original:  where δ i,j is the Kronecker delta and sparsemaxis an alternative to the softmax operator (Laha et al., 2016; Martins and Astudillo, 2016), which ensures the sparsity of the attention coefﬁcients and encourages the assignment of each node to a single centroid.
Revised: (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 deﬁned by: [Equation] p

Original: This alleviates the need for entropy minimization as done by DiffPool.
Revised: [Equation] 2 which corresponds to the euclidean projection of z onto the probability simplex ∆ K [Equation] =

Revised: [Equation]

Revised: p

Revised: [Equation]

Revised: p = 1 , p ≥ 0 (cid:9) .

Revised: The sparsemax operator ensures the sparsity of the attention coefﬁcients and encourages the assignment of each node to a single centroid.

Revised: It further alleviates the need for entropy minimization as done in DiffPool.

Original: Eq.
Revised: Eq.

Original: (5) also prevents the selected centroid nodes from being assigned to other clusters.
Revised: (5) also prevents the selected centroid nodes from being assigned to other clusters.

Original: Moreover, notice the term β i that regularizes the value of the attention for each node.
Revised: Moreover, notice the term β i that regularizes the value of the attention for each node.

Original: We can deﬁne β i = {+
+}
Revised: We can deﬁne β i = {+
+}

Original: [Equation]

Original:  VC , where d i , V C is the shortest path distance between each node v i ∈ V F and centroids in V C .
Revised: d i , VC , where d i , V C is the shortest path distance between each node v i ∈ V F and centroids in V C .

Original: Although this regularization incurs a cost O ( | V | 2 | V C | ) , it will strengthen the afﬁnity to closer centroids and ensure the connectivity of the resulting pooled graph.
Revised: Although this regularization incurs a cost O ( | V | 2 | V C | ) , it will strengthen the afﬁnity to closer centroids and ensure the connectivity of the resulting pooled graph.

Original: Note that this regularization is considered as an hyper-parameter of the layer and can be turned off, or alternatively, the mapping can be restricted to centroids within a ﬁxed h − hop neighborhood of each node.
Revised: 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 ﬁxed h − hop neighborhood of each node.

Original: Finally, after C l is computed at layer l , the coarsened graph G ( l +1) =
Revised: Finally, after C l is computed at layer l , the coarsened graph G ( l +1) =

Original: (cid:104) V ( l +1) ,
Revised: (cid:104) V ( l +1) ,

Original: A ( l +1) , X ( l +1) (cid:105) is computed using Eq.
Revised: A ( l +1) , X ( l +1) (cid:105) is computed using Eq.

Original: (6), as in (Ying et al., 2018).
Revised: (6), as in (Ying et al., 2018).

Original: 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.
Revised: 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.

Original: [Equation]
Revised: [Equation]

Original: This process can be repeated by feeding the new graph G ( l +1) into another GNN layer.
Revised: This process can be repeated by feeding the new graph G ( l +1) into another GNN layer.

Original: 3. P ROPERTIES OF THE L A P OOL M ETHOD
Revised: 3. P ROPERTIES OF THE L A P OOL M ETHOD

Original: 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 vertices set. Information PreservationThe centroid selection and the sparse distance-regularized node mapping ensure an appropriate segmentation of the graph and prevent the coarsened graph from being fully connected.
Revised: Preservation of Structural Information

Original: Although LaPool enforces a hierarchical structure on the graph, it still preserves the information of the original graph after a global sum-pooling (GSUM-Pool). Theorem 1.
Revised: 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.

Original: Deﬁne, the structure-aware feature content of a graph G = (cid:104) V, A, X (cid:105) as I ( G ) =
Revised: Substructure Identiﬁcation By construction, the soft assignment of nodes to centroids clusters existing substructures of the graph together, thus identifying important subgraphs according to the classiﬁcation task.

Original: AX .
Revised: By controlling for differences in signal variations within neighbourhoods, 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 ﬂexibility of determining the clustering dynamically, when training graphs sequentially, or statically when performing batch training.

Original: Ignoring the feature update performed by LaPool, GSUM-Pool( I ( LaPool ( G )) ) = GSUMPool( I ( G ) ). [Equation]
Revised: 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.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original:  R ESULTS AND D ISCUSSION
Revised:  R ESULTS AND D ISCUSSION

Original: A fundamental objective of LaPool is to learn an interpretable representation of sparse graphs, notably molecular substructures.
Revised: A fundamental objective of LaPool is to learn an interpretable representation of sparse graphs, notably molecular substructures.

Original: 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.
Revised: 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.

Original: This implies that GNNs should be able to identify semantically important substructure components from molecular graphs, and eventually reconstruct these graphs from such components.
Revised: This implies that GNNs should be able to identify semantically important substructure components from molecular graphs, and eventually reconstruct these graphs from such components.

Original: This stems from the intuition that molecular validity and functional properties derive more from chemical fragments than individual atoms.
Revised: This stems from the intuition that molecular validity and functional properties derive more from chemical fragments than individual atoms.

Original: 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:
Revised: 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:

Original: • LaPool’s consideration of semantically important information such as node distance translates to improved performance on molecular understanding and molecular activity prediction tasks.
Revised: • LaPool’s consideration of semantically important information such as node distance translates to improved performance on molecular understanding and molecular activity prediction tasks.

Original: • Visualization of LaPool’s behaviour at the pooling layer demonstrates its ability to identify coherent and meaningful molecular substructures.
Revised: • Visualization of LaPool’s behaviour at the pooling layer demonstrates its ability to identify coherent and meaningful molecular substructures.

Original: • The hierarchical representation enforced by LaPool, which preserves the original graph structure improves model interpretability.
Revised: • The hierarchical representation enforced by LaPool, which preserves the original graph structure improves model interpretability.

Original: • Learning meaningful substructures can be leveraged to construct a generative model which leads to more realistic and feasible molecules.
Revised: • Learning meaningful substructures can be leveraged to construct a generative model which leads to more realistic and feasible molecules.

Original: 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 ﬁlters each before the optional pooling layer, followed by 2 GNNs with 64 ﬁlters each and two fully connected layers including the output layer.
Revised: 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 ﬁlters each before the optional pooling layer, followed by 2 GNNs with 64 ﬁlters each and two fully connected layers including the output layer.

Original: Detailed information on architectural tuning and pooling-speciﬁc hyper-parameter search are provided in Supplemental Section A.
Revised: Detailed information on architectural tuning and pooling-speciﬁc hyper-parameter search are provided in Supplemental Section A.

Original: 4. B ENCHMARK ON M OLECULAR GRAPH UNDERSTANDING
Revised: 4. B ENCHMARK ON M OLECULAR GRAPH UNDERSTANDING

Original: While DiffPool and Graph U-Net models  outperform standard graph convolution networks on several graph benchmark datasets (Ying et al., 2018; Gao and Ji, 2018), we expect them to be less effective at identifying important substructures on molecular graphs as they do not explicitly consider structural relationships.
Revised: 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.

Revised: 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.

Original: We demonstrate this empirically by extracting known molecular substructure information from publicly available molecular datasets and evaluating performance in identifying these structures.
Revised: We demonstrate this empirically by extracting known molecular substructure information from publicly available molecular datasets and evaluating performance in identifying these structures.

Original: 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.
Revised: 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.

Original: As shown in Table [Table] the capture of structural relationships translates into superior performance of
Revised: As shown in Table 1, the capture of structural relationships translates into superior performance of

Original: LaPool, as measured across standard metrics on various substructure prediction tasks.
Revised: LaPool, as measured across standard metrics on various substructure prediction tasks.

Original: Indeed, we ﬁnd 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).
Revised: Indeed, we ﬁnd 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).

Original: In particular, on the harder  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.[Table]
Revised: 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.

Original: [Table]
Revised: [Table]

Original:  E XPERIMENTS ON STANDARD GRAPH CLASSIFICATION BENCHMARKS
Revised: 4. E XPERIMENTS ON STANDARD GRAPH CLASSIFICATION BENCHMARKS

Original: 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).
Revised: In addition to evaluating molecular structural understanding of the pooling models [Table] we benchmark our model on molecular toxicity prediction using the TOX21 dataset (Council et al., 2007).

Original: 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
Revised: 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

Original: Supplemental section C for dataset statistics).
Revised: Supplemental section C for dataset statistics).

Original: For TOX21, we report the test ROC-AUC averaged over 5-folds (following the 80-10-10 split proportion used in Wu et al.
Revised: For TOX [Table] we report the test ROC-AUC averaged over 5-folds (following the 80-10-10 split proportion used in Wu et al.

Original: (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.
Revised: (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.

Original: 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.
Revised: As shown in Table [Table] 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.

Original: In particular, on the PROTEINS dataset LaPool achieved an accuracy of 83 . 83 , representing a signiﬁcant gap relative to DiffPool, its closest competitor, at 77 . 25 .
Revised: In particular, on the PROTEINS dataset LaPool achieved an accuracy of 83 . 83 , representing a signiﬁcant gap relative to DiffPool, its closest competitor, at 77 . 25 .

Original: This suggests that the LaPool method is not restricted to molecular data but has broad applicability, especially in the context of sparse graph classiﬁcation.
Revised: This suggests that the LaPool method is not restricted to molecular data but has broad applicability, especially in the context of sparse graph classiﬁcation.

Original: [Table]
Revised: [Table]

Original: 4. M OLECULAR G ENERATION
Revised: 4. M OLECULAR G ENERATION

Original: We aim to showcase LaPool’s utility in drug discovery by demonstrating that it can be leveraged to generate molecules.
Revised: We aim to showcase LaPool’s utility in drug discovery by demonstrating that it can be leveraged to generate molecules.

Original: In previous work, GANs and VAEs were used to generate either string or graph representations of molecules.
Revised: In previous work, GANs and VAEs were used to generate either string or graph representations of molecules.

Original: 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
Revised: 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

Original: Supplemental Material).
Revised: Supplemental Material).

Original: For the encoder, we use a similar network architecture as in our supervised experiments.
Revised: For the encoder, we use a similar network architecture as in our supervised experiments.

Original: The decoder and discriminator are simple MLPs, with complete architectural details provided in Supplemental Section A.4.
Revised: The decoder and discriminator are simple MLPs, with complete architectural details provided in Supplemental Section A.4.

Original: Although the encoder is permutation invariant, the decoding process may not be.
Revised: Although the encoder is permutation invariant, the decoding process may not be.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: We ﬁnd that such a formulation improves the reconstruction loss and increases the ratio of valid molecules generated.
Revised: We ﬁnd that such a formulation improves the reconstruction loss and increases the ratio of valid molecules generated.

Original: Dataset and Baseline Models
Revised: Dataset and Baseline Models

Original: 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
Revised: 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

Original: QM9 molecular dataset (Ramakrishnan et al., 2014).
Revised: QM9 molecular dataset (Ramakrishnan et al., 2014).

Original: This dataset contains 133,885 small drug-like organic compounds with up to 9 heavy atoms (C, O, N, F).
Revised: This dataset contains 133,885 small drug-like organic compounds with up to 9 heavy atoms (C, O, N, F).

Original: 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).
Revised: 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).

Original: Our results are also compared to previous results on the same dataset, including Grammar-VAE, GraphVAE, and MolGAN.
Revised: Our results are also compared to previous results on the same dataset, including Grammar-VAE, GraphVAE, and MolGAN.

Original: Evaluation Metrics We measure the performance of the generative model using metrics standard in the ﬁeld: 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).
Revised: Evaluation Metrics We measure the performance of the generative model using metrics standard in the ﬁeld: 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).

Original: All metrics were computed on a set of 10,000 generated molecules.
Revised: All metrics were computed on a set of 10,000 generated molecules.

Original: [Table]
Revised: [Table]

Original: 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.
Revised: 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.

Original: Although MolGAN performed best on the novelty metric, it has among the lowest percentage of unique molecules.
Revised: Although MolGAN performed best on the novelty metric, it has among the lowest percentage of unique molecules.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We therefore conclude that the pooling performed by LaPool can improve molecular graph representation, which is crucial in a generative setting.
Revised: We therefore conclude that the pooling performed by LaPool can improve molecular graph representation, which is crucial in a generative setting.

Original: 4. I MPROVED I NTERPRETABILITY
Revised: 4. I MPROVED I NTERPRETABILITY

Original: To better understand the insights provided by LaPool, we conduct model interpretability experiments on molecular fragment prediction.
Revised: To better understand the insights provided by LaPool, we conduct model interpretability experiments on molecular fragment prediction.

Original: 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).
Revised: 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).

Original: This explains our focus on fragment prediction since in that setting an “interpretable model” that achieves high performance would need to ﬁrst understand the graph structure and the relationship between nodes. 
Revised: This explains our focus on fragment prediction since in that setting an “interpretable model” that achieves high performance would need to ﬁrst understand the graph structure and the relationship between nodes. [Figure]

Original: Wecompare LaPool to DiffPool, as it is the most similar method, and also because the node sampling performed by Graph U-net ignores the graph structure, typically disconnecting it. We ﬁrst investigate the behavior of each network by analyzing the clustering made at the pooling layer level. 
Revised: We ﬁrst 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.

Original: Indeed, 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.
Revised: 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.

Original: While deﬁning 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.
Revised: While deﬁning 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.

Original: 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.
Revised: 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.

Original: Indeed, the data-driven dynamic segmentation it performed is akin to chemical fragmentation (Gordon et al., 2011).
Revised: Indeed, the data-driven dynamic segmentation it performed is akin to chemical fragmentation (Gordon et al., 2011).

Original: In contrast,
Revised: In contrast,

Original: DiffPool’s cluster assignment is more uniform across the graph, leading to densely connected[Figure] c d e coarsened graphs which are less interpretable from a chemical viewpoint.
Revised: DiffPool’s cluster assignment is more uniform across the graph, leading to densely connected coarsened graphs which are less interpretable from a chemical viewpoint.

Revised: In particular, it fails in the presence of molecular symmetry, as it encourages the mapping of nodes with similar features to the same clusters.

Original: Example (c)  shows how DiffPool creates a fully connected graph from an originally disconnected graph, and example (b)  shows that symmetric elements, despite being far from each other, are assigned identically.
Revised: 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.

Original:  Such failures are not present when using LaPool.
Revised: 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.

Original: 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.
Revised: 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.

Original: However, this behavior is inherent to most DiffPool samples since the ﬁxed number of clusters and the inability to consider node distance cannot account for the diversity of molecular datasets.
Revised: However, this behavior is inherent to most DiffPool samples since the ﬁxed number of clusters and the inability to consider node distance cannot account for the diversity of molecular datasets.

Original: Cl Cl
Revised: Cl Cl

Original: Cl Cl
Revised: Cl Cl

Original: Cl Cl Cl
Revised: Cl Cl Cl

Original: O
Revised: O

Original: O HO
Revised: O HO

Original: O
Revised: O

Original: O H H H H
Revised: O H H H H

Original: C
Revised: C

Original: C
Revised: C

Original: O2 O
Revised: O2 O

Original: C4 C5 C6 C
Revised: C4 C5 C6 C

Original: C
Revised: C

Original: C9 C
Revised: C9 C

Original: C11 C
Revised: C11 C

Original: C C
Revised: C C

Original: O
Revised: O

Original: O16 C17 C18 C
Revised: O16 C17 C18 C

Original: C
Revised: C

Original: C21 C
Revised: C21 C

Original: C23 C24 C C26 C27 H28 H
Revised: C23 C24 C25 C26 C27 H28 H

Original: H30 H31 score = 0.
Revised: H30 H31 score = 0.

Original: C
Revised: C

Original: C
Revised: C

Original: O2 O
Revised: O2 O

Original: C4 C5 C6 C
Revised: C4 C5 C6 C

Original: C
Revised: C

Original: C9 C
Revised: C9 C

Original: C11 C
Revised: C11 C

Original: C C
Revised: C C

Original: O
Revised: O

Original: O16 C17 C18 C
Revised: O16 C17 C18 C

Original: C
Revised: C

Original: C21 C
Revised: C21 C

Original: C23 C24 C C26 C27 H 28 H
Revised: C23 C24 C25 C26 C27 H 28 H

Original: H30 H31 score = 0.
Revised: H30 H31 score = 0.

Original: C

Original: C

Original: O2 OC4 C5 C6 C C C9 C
Revised: O

Original: C11 C

Original: C

Original: C

Original: O
Revised: O

Original: O16 C17 C18 C

Original: C
Revised: C

Original: C21 CC23 C24 C C26 C27 H28 H H30 score = 1.
Revised: C

Original: O O C C
Revised: C

Original: C C O
Revised: C3 O

Original: O5 C6 C
Revised: O5 C6 C

Original: C8 C score = 
Revised: C8 C score = 0.

Original: [Table] C C C
Revised: C

Original: CO O5 C6 C C8 C score =
Revised: C

Original: [Table]

Original: C
Revised: C

Original: C C C O
Revised: C3 O

Original: O5 C6 C
Revised: O5 C6 C

Original: C8 C score = 
Revised: C8 C score = 0.

Original: [Table]

Original: O O
Revised: O O

Original: O O N
Revised: O O N

Original: N
Revised: N

Original: O O O O
Revised: O O O O

Original: O
Revised: O

Original: O
Revised: O

Original: H
Revised: H

Original: H H
Revised: H H

Original: Cl Cl
Revised: Cl Cl

Original: Cl Cl
Revised: Cl Cl

Original: Cl Cl

Revised: Cl Cl

Original: Cl
Revised: Cl

Original: O
Revised: O

Original: C C C10 C
Revised: C8 C9 C10 C

Original: C
Revised: C

Original: C C C C16 C 
Revised: C13 C14 C15 C16 C score = 0.

Original: [Table] C[Table]
Revised: C0 C

Original: C2 C3 C4 C
Revised: C2 C3 C4 C

Original: C6 C7  C
Revised: C6 C7 C8 C

Original: C9 C
Revised: C

Original: C11 C12 C13 C
Revised: C11 C12 C13 C

Original: C15 C16 O17 O
Revised: C15 C16 O17 O

Original: O19 O
Revised: O19 O

Original: C21 C22 C
Revised: C21 C22 C

Original: C24 C25 C
Revised: C24 C25 C

Original: C27 C28 N
Revised: C27 C28 N

Original: N [Table] = 
Revised: N score = 1.

Original: [Table]

Original: [Table]

Original: C C
Revised: C C

Original: C
Revised: C

Original: O3 O
Revised: O3 O

Original: O
Revised: O

Original: C C7 C
Revised: C6 C7 C

Original: C C
Revised: C C

Original: C
Revised: C

Original: C12 C13 C
Revised: C12 C13 C

Original: O
Revised: O

Original: O O
Revised: O O

Original: C18 C
Revised: C18 C

Original: C C21 C22 C
Revised: C20 C21 C22 C

Original: C C25 C
Revised: C24 C25 C

Original: C27 C C H
Revised: C27 C C29 H

Original: H
Revised: H

Original: H score = [Table]
Revised: H score = 0.

Original: Cl
Revised: Cl

Original: Cl Cl
Revised: Cl Cl

Original: Cl Cl
Revised: Cl Cl

Original: Cl
 Cl
Revised: Cl
 Cl

Original: O
Revised: O

Original: C

Original: C

Original: C10 C

Original: C

Original: C C C C16 C score = 
Revised: C8 C9 C10 C11 C12 C13 C14 C15 C16 C score = 0.

Original: [Table]

Original: C0 C
Revised: C0 C

Original: C2 C3 C4 C
Revised: C2 C3 C4 C

Original: C6 C7  C
Revised: C6 C7 C8 C

Original: C9 C
Revised: C

Original: C11 C12 C13 C
Revised: C11 C12 C13 C

Original: C15 C16 O17 O
Revised: C15 C16 O17 O

Original: O19 O
Revised: O19 O

Original: C21 C22 C
Revised: C21 C22 C

Original: C24 C25 C
Revised: C24 C25 C

Original: C27 C28 N
Revised: C27 C28 N

Original: N score = 
Revised: N score = 1.

Original: [Table]

Original: C C
Revised: C C

Original: C
Revised: C

Original: O3 O
Revised: O3 O

Original: O
Revised: O

Original: C C7 C
Revised: C6 C7 C

Original: C C
Revised: C C

Original: C
Revised: C

Original: C12 C13 C
Revised: C12 C13 C

Original: O
Revised: O

Original: O O
Revised: O O

Original: C18 C
Revised: C18 C

Original: C C21 C22 C
Revised: C20 C21 C22 C

Original: C  C25 C
Revised: C 24 C25 C

Original: C27 C C H
Revised: C27 C C29 H

Original: H
Revised: H

Original: H score = 0.
Revised: H score = 0.

Original: Cl
Revised: GIN C C O2 O

Revised: C4 C5 C6 C

Revised: C

Revised: C9 C

Original: Cl
Revised: C11 C C C O

Revised: O16 C17 C18 C

Revised: C

Revised: C21 C

Revised: C23 C24 C25 C26 C27 H28 H

Revised: H30 score = 1.

Original: Cl
Revised: C C C C3 O

Revised: O5 C6 C

Revised: C8 C score = 0.

Original: Cl 
Revised: Cl Cl Cl Cl

Original: Cl {+Cl
 Cl O+}
Revised: Cl {+Cl
 Cl O+}

Revised: C8 C9 C10 C11 C12 C13 C14 C15 C16 C score = 1.

Revised: C0 C

Revised: C2 C3 C4 C

Revised: C6 C7 C8 C

Revised: C

Revised: C11 C12 C13 C

Revised: C15 C16 O17 O

Revised: O19 O

Revised: C21 C22 C

Revised: C24 C25 C

Revised: C27 C28 N

Revised: N score = 1.

Original: [-Cl
-]
Revised: [-Cl
-]C C C O3 O

Revised: O

Revised: C6 C7 C

Original: Cl
Revised: C C C C12 C13 C

Revised: O

Revised: O

Revised: O

Revised: C18 C

Revised: C20 C21 C22 C

Revised: C24 C25 C

Revised: C27 C

Revised: C29 H

Revised: H

Revised: H score = 0.

Revised: LAPOOL

Revised: C

Revised: C

Revised: O2 O

Revised: C4 C5 C6 C

Revised: C

Revised: C9 C

Revised: C11 C

Revised: C

Revised: C

Original: O
Revised: O

Revised: O16 C17 C18 C

Revised: C

Revised: C21 C

Revised: C23 C24 C25 C26C27 H28 H

Revised: H30 H

Revised: C

Revised: C

Revised: C

Revised: C3 O

Revised: O5 C6 C

Original:  C 
Revised: C8 C score = 0. Cl Cl Cl

Original: C{+Cl Cl Cl
 Cl+}
Revised: {+Cl Cl Cl
 Cl+}

Original:  C10 C
Revised: O C8 C9 C10 C

Original: C
Revised: C

Revised: C13 C14 C15 C16 C score = 0.

Original: C C
Revised: C C

Original: C
Revised: C

Original: C16 C score = 1.
Revised: O3 O O

Original: C0 C
Revised: C6 C7 C C C C

Revised: C12 C13 C

Revised: O

Revised: O

Original: C2 C3 C4 C
Revised: O C18 C C20 C21 C22 C C24 C25 C

Original: C6 C7 C
Revised: C27 C

Original: C9 C
Revised: C29 H

Original: C11 C12 C13 C
Revised: H

Original: C15 C16 O17 O O19 O
Revised: H score = 0. C0 C

Original: C21 C22 C C24 C25 C
Revised: C2 C3 C4 C C6 C7 C8 C

Original: C27 C28 N
Revised: C

Original: N score = 1. C C C
Revised: C11 C12 C13 C

Original: O3 O OC C7 C
Revised: C15 C16 O17 O

Original: C

Original: C

Original: C

Original: C12 C13 C

Original: O O O
Revised: O19 O

Original: C18 C C C21 C22 C
Revised: C21 C22 C

Original: C C25 C
Revised: C24 C25 C

Original: C27 C C H
Revised: C27 C28 N

Original: H
Revised: N

Original: H score = 0. GIN
Revised: GRAPH U-Net score = 1. score = 1.

Original: DIFFPOOL
Revised: DIFFPOOL

Original: [Figure]
Revised: GROUND TRUTH

Original: 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.
Revised: 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.

Original: We design a simple experiment in which we predict the presence of eitherEpoxide,
Revised: We design a simple experiment in which we predict the presence of either

Original:  Thioepoxide, or Aziridine substructures (denoted by the molecular pattern “C1[O,S,N]C1”), that are indicative of molecular toxicity.
Revised: Epoxide, Thioepoxide, or Aziridine substructures (denoted by the molecular pattern “C1[O,S,N]C1”), that are indicative of molecular toxicity.

Original: Interpretability is therefore deﬁned 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.
Revised: Interpretability is therefore deﬁned 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.

Original: Similar to (Pope et al., 2019), we adapt an existing explainability method for CNNs to GNNs.
Revised: Similar to (Pope et al., 2019), we adapt an existing explainability method for CNNs to GNNs.

Original: Speciﬁcally, 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 D for discussion and alternate approach).
Revised: Speciﬁcally, 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).

Original: Next, we derive an importance score for each node using the L1-norm of the feature attribution map for the node.
Revised: Next, we derive an importance score for each node using the L1-norm of the feature attribution map for the node.

Original: 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 ﬁnd LaPool to more robustly identify the salient structure, resulting in improved overall interpretability. 
Revised: 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 ﬁnd 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.

Revised: Such performance is a direct result of using a cluster size large enough to cover the toxic fragment size.

Original: [Table]
Revised: [Table]

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we have proposed LaPool, a differentiable and robust pooling operator for molecular  graphs that considers both node information and graph structure.
Revised: 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.

Revised: In doing so, we have proposed a method which is able to identify important substructures of a graph by leveraging the graph Laplacian.

Revised: 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.

Original: By 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.
Revised: 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.

Original: We discussed the performance of LaPool relative to existing graph pooling layers and demonstrated on both molecular graph classiﬁcation and generation benchmarks that LaPool outperforms existing graph pooling modules and produces more interpretable results.
Revised: We discussed the performance of LaPool relative to existing graph pooling layers and demonstrated on both molecular graph classiﬁcation and generation benchmarks that LaPool outperforms existing graph pooling modules and produces more interpretable results.

Original: 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.
Revised: 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.

Original: Finally, we show that although LaPool was designed for molecular graphs, it generalizes well to other graph types.
Revised: Finally, we show that although LaPool was designed for molecular graphs, it generalizes well to other graph types.

Original: In future work, we aim to investigate how additional sources of information such as edge features could be incorporated into the graph pooling process
Revised: In future work, we aim to investigate how additional sources of information such as edge features could be incorporated into the graph pooling process

Original: A N ETWORK A RCHITECTURE AND TRAINING PROCEDURE
Revised: A N ETWORK A RCHITECTURE AND TRAINING PROCEDURE

Original: Below, we describe the network architecture and the training procedure used for both supervised and generative experiments.
Revised: Below, we describe the network architecture and the training procedure used for both supervised and generative experiments.

Original: A. E DGE ATTRIBUTES
Revised: A. E DGE ATTRIBUTES

Original: Part of the work presented assumes the absence of edge attributes in the graphs.
Revised: Part of the work presented assumes the absence of edge attributes in the graphs.

Original: However, in molecular graphs, the nature of a bond between two atoms plays an important role regarding activity and property.
Revised: However, in molecular graphs, the nature of a bond between two atoms plays an important role regarding activity and property.

Original: As such, edge types should be considered, especially in generative models.
Revised: As such, edge types should be considered, especially in generative models.

Original: To consider this, we add to our network an initial Edge-GC layer described in the following.
Revised: To consider this, we add to our network an initial Edge-GC layer described in the following.

Original: Let G =
Revised: Let G =

Original: (cid:104) V, E, X (cid:105) be an undirected molecular graph, such that E
Revised: (cid:104) V, E, X (cid:105) be an undirected molecular graph, such that E

Original: =
Revised: =

Original: [ E 1 , . . .
Revised: [ E 1 , . . .

Original: E k ] ∈ { 0 , 1 } e × n × n where n is the number of nodes in the graph and e is the number of possible edge.
Revised: E k ] ∈ { 0 , 1 } e × n × n where n is the number of nodes in the graph and e is the number of possible edge.

Original: We have that
Revised: We have that

Original: [Equation]
Revised: [Equation]

Original: where A is the adjacency matrix of the graph.
Revised: where A is the adjacency matrix of the graph.

Original: The Edge GC layer is simply deﬁned as :
Revised: The Edge GC layer is simply deﬁned as :

Original: [Equation]
Revised: [Equation]

Original: where (cid:107) 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.
Revised: where (cid:107) 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.

Original: A new graph deﬁned as
Revised: A new graph deﬁned as

Original: G (cid:48)
Revised: G (cid:48)

Original: = (cid:104) V, A, Y (cid:105) can then be feed into the subsequent layers of the network.
Revised: = (cid:104) V, A, Y (cid:105) can then be feed into the subsequent layers of the network.

Original: A. A TOM AND EDGE FEATURES
Revised: A. A TOM AND EDGE FEATURES

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: A. S UPERVISED EXPERIMENTS
Revised: A. S UPERVISED EXPERIMENTS

Original: 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.
Revised: 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.

Original: This is further followed by one fully connected layers (128) with batch normalization and ReLU activation, ﬁnalized by a linear output layer with appropriate activation for the task readouts.
Revised: This is further followed by one fully connected layers (128) with batch normalization and ReLU activation, ﬁnalized by a linear output layer with appropriate activation for the task readouts.

Original: Notice that we used one pooling layer, since no noticeable improvement was observed when using more in our experimental setting.
Revised: Notice that we used one pooling layer, since no noticeable improvement was observed when using more in our experimental setting.

Original: For DiffPool, we performed a hyperparameter search to ﬁnd the optimal number of clusters (12.5%, 25%, 50% of the maximum number of nodes in the batch (Ying et al., 2018)).
Revised: For DiffPool, we performed a hyperparameter search to ﬁnd the optimal number of clusters (12.5%, 25%, 50% of the maximum number of nodes in the batch (Ying et al., 2018)).

Original: A similar search is also performed for the Graph-Unet pooling layer.
Revised: A similar search is also performed for the Graph U-net pooling layer.

Original: For LaPool, 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.
Revised: 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.

Original: The grid search was performed for k ∈ { 1 , 2 , 3 } .
Revised: The grid search was performed for k ∈ { 1 , 2 , 3 } .

Original: For the supervised experiments, we use a batch size of 64 and train the networks for 100 epochs, with early stopping.
Revised: For the supervised experiments, we use a batch size of 64 and train the networks for 100 epochs, with early stopping.

Original: A. G ENERATIVE MODELS A.4. WAE MODEL
Revised: A. G ENERATIVE MODELS A.4. WAE MODEL

Original: We use a Wasserstein Auto-Encoder (WAE) as our generative model (see Figure S1.
Revised: We use a Wasserstein Auto-Encoder (WAE) as our generative model (see Figure S1.

Original: The WAE minimizes a penalized form of the Wasserstein distance between a model distribution and a target
Revised: The WAE minimizes a penalized form of the Wasserstein distance between a model distribution and a target

Original: G N N
Revised: G N N

Original: L a y e r
Revised: L a y e r

Original: E - G
Revised: E - G

Original: NN L a y e r
Revised: NN L a y e r

Original: G r a p h
Revised: G r a p h

Original: Poo li n g
Revised: Poo li n g

Original: G - S u m Poo l
Revised: G - S u m Poo l

Original: F C L a y e r
Revised: F C L a y e r

Original: O O
Revised: O O

Original: NH
Revised: NH

Original: S E
Revised: S E

Original: G
Revised: G

Original: O O
Revised: O O

Original: NH 0/
Revised: NH 0/

Original: D
Revised: D

Original: G N N
Revised: G N N

Original: L a y e r ... ll
Revised: L a y e r ... ll

Original: [Figure]
Revised: [Figure]

Original: distribution, and has been shown to improve learning stability.
Revised: distribution, and has been shown to improve learning stability.

Original: As described in (Tolstikhin et al., 2017), we aim to minimize the following objective:
Revised: As described in (Tolstikhin et al., 2017), we aim to minimize the following objective:

Original: [Equation]
Revised: [Equation]

Original: where Q is any nonparametric set of probabilistic encoders, D Z is the Jensen-Shannon divergence between the learned latent distribution Q Z and prior P Z , and λ > 0 is a hyperparameter.
Revised: where Q is any nonparametric set of probabilistic encoders, D Z is the Jensen-Shannon divergence between the learned latent distribution Q Z and prior P Z , and λ > 0 is a hyperparameter.

Original: D Z is estimated using adversarial training (discriminator).
Revised: D Z is estimated using adversarial training (discriminator).

Original: 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 = { G 1 , · · · , G n } .
Revised: 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 = { G 1 , · · · , G n } .

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: The output of those layers is used as shared embedding for two networks:
Revised: The output of those layers is used as shared embedding for two networks:

Original: one predicting the upper triangular entries of the edge tensor, and the second predicting the node features tensor.
Revised: one predicting the upper triangular entries of the edge tensor, and the second predicting the node features tensor.

Original: This results in faster convergence.
Revised: This results in faster convergence.

Original: For the discriminator, we use a simple MLP that predicts whether the latent code comes from a normal prior distribution z N (0 , 1) .
Revised: For the discriminator, we use a simple MLP that predicts whether the latent code comes from a normal prior distribution z N (0 , 1) .

Original: This MLP is constituted by two stacked FCLs (64, 32) followed by an output layer with sigmoid activation.
Revised: This MLP is constituted by two stacked FCLs (64, 32) followed by an output layer with sigmoid activation.

Original: As in (Kadurin et al., 2017), we do not use batch-normalization, since it resulted in a mismatch between the discriminator and the generator.
Revised: As in (Kadurin et al., 2017), we do not use batch-normalization, since it resulted in a mismatch between the discriminator and the generator.

Original: All models use the same basic generative architecture, with the only difference being the presence of a pooling-layer and its associated parameters.
Revised: All models use the same basic generative architecture, with the only difference being the presence of a pooling-layer and its associated parameters.

Original: For DiffPool, we ﬁxed the number of cluster to three, while for LaPool, we use the distance-based regularization.
Revised: For DiffPool, we ﬁxed the number of cluster to three, while for LaPool, we use the distance-based regularization.

Original: A.4. R ECONSTRUCTION LOSS
Revised: A.4. R ECONSTRUCTION LOSS

Original: For each input molecular graph G = (cid:104) V, E, X (cid:105) , the decoder reconstruct a graph ˜ G = (cid:104) ˜ V , ˜ E, ˜ X (cid:105) .
Revised: For each input molecular graph G = (cid:104) V, E, X (cid:105) , the decoder reconstruct a graph ˜ G = (cid:104) ˜ V , ˜ E, ˜ X (cid:105) .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In (Simonovsky and
Revised: In (Simonovsky and

Original: Komodakis), the authors use an expensive graph matching procedure to overcome that limitation.
Revised: Komodakis), the authors use an expensive graph matching procedure to overcome that limitation.

Original: We argue that it sufﬁces to compute the reconstruction loss on γ ( G ) and γ ( ˜ G ) , where γ is a permutation invariant embedding function.
Revised: We argue that it sufﬁces to compute the reconstruction loss on γ ( G ) and γ ( ˜ G ) , where γ is a permutation invariant embedding function.

Original: As a heuristic, we used a Graph Isomorphism Network (GIN), with weights ﬁxed to 1, in order to approximate the Weisfeiler-Lehman graph isomorphism test (see (Xu et al., 2018) for more details).
Revised: As a heuristic, we used a Graph Isomorphism Network (GIN), with weights ﬁxed to 1, in order to approximate the Weisfeiler-Lehman graph isomorphism test (see (Xu et al., 2018) for more details).

Original: In particular, we use an edge-aware GIN layer (see section A.1) to embed both G and ˜ G .
Revised: In particular, we use an edge-aware GIN layer (see section A.1) to embed both G and ˜ G .

Original: The reconstruction loss is then deﬁned as:
Revised: The reconstruction loss is then deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: A.4. T RAINING PROCEDURE
Revised: A.4. T RAINING PROCEDURE

Original: The QM9 dataset was split into a train (60%), valid (20%) and a hold-out test dataset (20%).
Revised: The QM9 dataset was split into a train (60%), valid (20%) and a hold-out test dataset (20%).

Original: Note that only 25% of the training set is sampled during each epoch (batch size 32).
Revised: Note that only 25% of the training set is sampled during each epoch (batch size 32).

Original: 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 1 e − 4 for the generator and 1 e −for the discriminator.
Revised: 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 1 e − 4 for the generator and 1 e −for the discriminator.

Original: During training, we slowly reduce the learning rate by a factor of 0.5, for the generator, on plateau.
Revised: During training, we slowly reduce the learning rate by a factor of 0.5, for the generator, on plateau.

Original: 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.
Revised: 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.

Original: B G ENERATED MOLECULES
Revised: B G ENERATED MOLECULES

Original: Below, we highlight a few molecules generated by WAE-LaP on the QM9 dataset.
Revised: Below, we highlight a few molecules generated by WAE-LaP on the QM9 dataset.

Original: OH
Revised: OH

Original: NH
Revised: NH

Original: NH NH
Revised: NH NH

Original: HO
Revised: HO

Original: O
Revised: O

Original: NH
Revised: NH

Original: F HO
Revised: F HO

Original: H 2 N
Revised: H 2 N

Original: NH
Revised: NH

Original: N NH
Revised: N NH

Original: NH 2 HO
Revised: NH 2 HO

Original: NH NH
Revised: NH NH

Original: OH
Revised: OH

Original: NH
Revised: NH

Original: NH O
Revised: NH O

Original: OH
Revised: OH

Original: O
Revised: O

Original: NH
Revised: NH

Original: O O
Revised: O O

Original: OH
Revised: OH

Original: NH
Revised: NH

Original: HO
Revised: HO

Original: NH NH
Revised: NH NH

Original: O
Revised: O

Original: OH
Revised: OH

Original: HO
Revised: HO

Original: NH
Revised: NH

Original: HO
Revised: HO

Original: O
Revised: O

Original: OH N
Revised: OH N

Original: NH
Revised: NH

Original: OH HO
Revised: OH HO

Original: NH
Revised: NH

Original: OH
Revised: OH

Original: OH HO
Revised: OH HO

Original: C D ATASET STATISTICS
Revised: C D ATASET STATISTICS

Original: TOX
Revised: TOX

Original: DD
Revised: DD

Original: PROTEINS
Revised: PROTEINS

Original: FRANKENSTEIN
Revised: FRANKENSTEIN

Original: [Table]
Revised: [Table]

Original: D N ODE IMPORTANCE INTERPRETABILITY SCORE
Revised: D N ODE IMPORTANCE INTERPRETABILITY SCORE

Original: In addition to assessing the quality of clustering performed by LaPool and DiffPool
Revised: In addition to assessing the quality of clustering performed by LaPool and DiffPool

Original: [Table]
Revised: [Table]

Original: we attempt to measure the interpretability of their predictions.
Revised: we attempt to measure the interpretability of their predictions.

Original: We consider a setting in which the goal is to predict the presence of either Epoxides, Thioepoxides or Aziridines substructures in molecular graphs.
Revised: We consider a setting in which the goal is to predict the presence of either Epoxides, Thioepoxides or Aziridines substructures in molecular graphs.

Original: These three fragments correspond to structural alerts that are often indicative of molecular toxicity.
Revised: These three fragments correspond to structural alerts that are often indicative of molecular toxicity.

Original: We attempt to identify the relational inductive bias used by each model during prediction.
Revised: We attempt to identify the relational inductive bias used by each model during prediction.

Original: In our setting, we deﬁne 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.
Revised: In our setting, we deﬁne 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.

Original: In other words, we expect the most important nodes for the model prediction to correspond to nodes that are part of the structural alerts.
Revised: In other words, we expect the most important nodes for the model prediction to correspond to nodes that are part of the structural alerts.

Original: We measure the importance of each atom toward the model prediction using the Integrated gradient method.
Revised: We measure the importance of each atom toward the model prediction using the Integrated gradient method.

Original: Brieﬂy, 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.
Revised: Brieﬂy, 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.

Original: 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.
Revised: 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.

Original: Finally, we compute the interpretability score using the Precision-Recall AUC between measured importance and ground truth which is deﬁned by a binary mask of nodes that are part of the structural alerts.
Revised: Finally, we compute the interpretability score using the Precision-Recall AUC between measured importance and ground truth which is deﬁned by a binary mask of nodes that are part of the structural alerts.

Original: The
Revised: The

Original: PR-AUC allows us to assess the node importance separation capacity of each model while taking imbalance into account.
Revised: PR-AUC allows us to assess the node importance separation capacity of each model while taking imbalance into account.

Original: We only focus on positive predictions for each model.
Revised: We only focus on positive predictions for each model.

Original: As an alternative to the Integrated Gradient, we also measure the interpretability score using Guided BackPropagation (see Table 6)
Revised: As an alternative to the Integrated Gradient, we also measure the interpretability score using Guided BackPropagation (see Table 6)

Original: [Table]
Revised: [Table]

Original: E S IGNAL PRESERVATION THROUGH L APLACIAN MAXIMA
Revised: E S IGNAL PRESERVATION THROUGH L APLACIAN MAXIMA

Original: 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).
Revised: 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).

Original: We measure the energy conservation after downsampling: δ E ( S ) =
Revised: We measure the energy conservation after downsampling: δ E ( S ) =

Original: E ( S ) − E
Revised: E ( S ) − E

Original: ( S down ) 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).
Revised: ( S down ) 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).

Original: The energy E S of a discrete signal y
Revised: The energy E S of a discrete signal y

Original: i is deﬁned in (11), and is similar to the energy of a wave in a physical system (without the constants).
Revised: i is deﬁned in (11), and is similar to the energy of a wave in a physical system (without the constants).

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: S
Revised: S

Original: i g n a l i n t e n s i t y
Revised: i g n a l i n t e n s i t y

Original: [Figure]
Revised: [Figure]

Original: 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.
Revised: 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.

Original: For the pooling methods, a linear interpolation is used to cover the same signal space before computing E S .
Revised: For the pooling methods, a linear interpolation is used to cover the same signal space before computing E S .

Original: As expected, the maxima Laplacian selection requires a fewer number of samples for signal reconstruction and energy preservaton.
Revised: As expected, the maxima Laplacian selection requires a fewer number of samples for signal reconstruction and energy preservaton.

Original: It also signiﬁcantly outperforms minima selection.
Revised: It also signiﬁcantly outperforms minima selection.