Memory-Based Graph Networks
Authors: Amir Hosein Khasahmadi
Kaveh Hassani
Parsa Moradi
Leo Lee
Quaid Morris
Venue:  (2020)
Keywords: Graph Neural Networks
Memory Networks
Hierarchial Graph Representation Learning

Abstract: Graph neural networks (GNNs) are a class of deep models that operate on data with arbitrary topology represented as graphs. We introduce an efficient memory layer for GNNs that can jointly learn node representations and coarsen the graph. We also introduce two new networks based on this layer: memory-based GNN (MemGNN) and graph memory network (GMN) that can learn hierarchical graph representations. The experimental results shows that the proposed models achieve state-of-the-art results in eight out of nine graph classification and regression benchmarks. We also show that the learned representations could correspond to chemical features in the molecule data.


Published as a conference paper at ICLR 2020
MEMORY-BASED GRAPH NETWORKS
Amir Hosein Khasahmadi 2 Kaveh Hassani", Parsa Moradi', Leo Leet:2, Quaid Morris1.2
l University of Toronto, Toronto, Canada
Vector Institute, Toronto, Canada
3 "Autodesk AI Lab, Toronto, Canada
Sharif University of Technology, Tehran, Iran kaven. aaassanieautooeeee quaid.morris www.rorrrittttttttoooooo
{amirhesein, 1 f toronto.. edu parsa. mordadieee.d....mooo shari Sharif.eet e.edu
ABSTRACT
Graph neural networks (GNNs) class of deep models that operate data with arbitrary topology represented as graphs. We introduce an efficient mem- ory layer for GNNs that can jointly learn node representations and coarsen the graph. We also introduce two new networks based on this layer: memory-based GNN (MemGNN) and graph memory network (GMN) that can learn hierarchical graph representations. The experimental results show that the proposed models achieve state-of-the-arr results in eight out of nine graph classification and regres- sion benchmarks. We also show that the learned representations could correspond to chemical features in the molecule data. Code and reference implementations are released at: https://github.. wwwptthub..oommnii
INTRODUCTION
Graph neural networks (GNNs) (Wu et al. 2019 Zhou et al. 2018 Zhang al. 2018) are class of deep models that operate on data with arbitrary topology represented as graphs such as social networks (Kipf & Welling 2016), knowledge graphs (Vivona & Hassani 2019), molecules (Duve- naud et al. 2015), point clouds (Hassani & Haley] 2019).. and robots (Wang al. 2019). Unlike regular-structuree structured inputs such as grids (c.g., images and volumetric data) and sequences (c.g.. speech and text), GNN inputs
variable-size graphs consisting of nodes and their interactions. GNNs such as gated GNN (GGNN) (Li et al.] 2015), message passing neural net- work (MPNN) (Gilmer al. 2017). graph convolutional network (GCN) (Kipf & Welling 2016), and graph attention network (GAT) (Velikovi et al., 2018) learn node representations through an iterative process of transferring, transforming, and aggregating the node representations from topo- logical neighbors. Each iteration expands the receptive field by one hop and after & iterations the nodes within & hops influence the node representations of one another. GNNs are shown learn better representations compared to random walks (Grover & Leskovec) 2016 Perozzi et al.) 2014). matrix factorization (Belkin & Niyogi 2002) Ou al. (016). kernel methods (Shervashidze al., Kriege et al. 2016), and probabilistic graphical models (Dai et al. 2016)).
permutation-invariae
2011
These models, however, cannot learn hierarchical representations as they do not exploit the compo- sitional nature of graphs. Recent work such as differentiable pooling (DiffPool) (Ying et al. 2018). TopKPool (Gao & Ji, (201)). and self-attention graph pooling (SAGPool) (Lee et al., 2019 intro- duce parametric graph pooling layers that allow GNNs to learn hierarchical graph representations by stacking interleaved layers of GNN and pooling layers. These layers cluster nodes in the latent space. The clusters may correspond communities in social network potent functional groups within a chemical dataset. Nevertheless, these models are not efficient as they require an iterative process of message passing after each pooling layer.
In this paper, we introduce memory layer for joint graph representation learning and graph coars- ening that consists of a multi-head array of memory keys and convolution operator aggregate the soft cluster assignments from different heads. The queries to memory layer are node repre- sentations from the previous layer and the outputs are the node representations of the coarsened graph. The memory layer does not explicitly require connectivity information and unlike GNNs
Work done during internship Autodesk Toronto AT Lah.

---

Published as a conference paper at ICLR 2020
relies on the global information rather than local topology. Hence, it does not struggle with over- smoothing problem (Xu et al. 2018; Li et al. 2018). These properties make memory layers more efficient and improves their performance. We also introduce two networks based on the proposed layer: memory-based GNN (MemGNN) and graph memory network (GMN). MemGNN consists of GNN that learns the initial node representations, and stack of memory layers that lear hi- erarchical representations to the global graph representation. GMN, on the other hand, learns the hierarchical representations purely based on memory layers and hence does not require message passing.
2 RELATED WORK
Memory augmented neural networks (MANNs) utilize external memory with differentiable read- write operators allowing them to explicitly access the past experiences and are shown to enhance reinforcement learning (Pritzel et al. 2017), meta learning (Santoro et al. (0166), few-shot learning (Vinyals et al. 2016), and multi-hop reasoning (Weston et al. 2015). Unlike RNNs, in which the memory is represented within their hidden states, the decoupled memory in MANNs allows them to store and retrieve longer term memories with less parameters. The memory can be implemented as: key-value memory such as neural episodic control (Pritzel et al., 2017) and product-key memory layers (Lample et al., 2019), or array-structured memory such as neural Turing machine (NTM) (Graves et al., 2014), prototypical networks (Snell al. 2017). memory networks (Weston al., 2013), and sparse access memory (SAM) (Rae et al. 2016). Our memory layer consists of a multi- head array of memory keys.
Graph neural networks (GNNs) mostly use message passing to learn node representations over graphs. GraphSAGE (Hamilton et al. (017) learns representations by sampling and aggregating neighbor nodes whereas GAT (Velikovi et al.) 2018) uses attention mechanism to aggregate repre- sentations from all neighbors. GCN extend the convolution operator to arbitrary topology. Spectral GCNs (Bruna et al. 2014 Defferrard et al., 2016; Kipf & Welling 2016) use spectral filters over graph Laplacian to define the convolution in the Fourier domain. These models are less efficient compared to spatial GCNs (Schlichtkrull et al. 2018 Ma et al., 2019) which directly define the convolution on graph patches centered on nodes. Our memory layer uses a feed-forward network to learn the node representations.
Graph pooling can be defined in global or hierarchical manners. In former, node representations are aggregated into a graph representation by readout layer implemented using arithmetic operators such as summation or averaging (Hamilton et al.] 2017: Kipf & Welling 2016) or set neural net- works such as Set2Set (Vinyals et | and SortPool (Morris et al., 2019). In latter, graphs are coarsened in each layer to capture the hierarchical structure. Efficient non- con-parametric methods such as clique pooling (Luzhnica et al. 2019), kNN pooling (Wang et al., 2018). and Graclus (Dhillon)
eet al., 2007) rely on topological information but are outperformed by parametric models such as edge contraction pooling (Diehl 2019).
DiffPool (Ying et al.) 2018)) trains two parallel GNNs to compute node representations and clus- ter assignments using a multi-term loss including classification, link prediction, and entropy losses, whereas Mincut pool (Bianchi et al., (219) trains sequence of GNN and an MLP using classifica- tion loss and the minimum cut objective. TopKPool (Cangea et al. 2018;) Gao & Ji 2019) computes node scores by learning projection vectors and dropping all the nodes except the top scoring nodes. SAGPool (Lee et al.] 2019) extends the TopKPool by using graph convolutions to take neighbor node features into account. We use clustering frreinngggrreyyyyy distribution to compute the attention scores between nodes and clusters.
3 METHOD
3.1 MEMORY LAYER
We define a memory layer M Ruidd ب R xd4: in layer as a parametric function that takes in ny query vectors of size di and generates n141 query vectors of size d1+1 such that 14+1 < ny. The input and output queries represent the node representations of the input graph and the coarsened graph, respectively. The memory layer learns jointly coarsen the input nodes,
2

---

Published as a conference paper at ICLR 2020
Queries
Memory Layer
Memory Layer
Input Graph
5000
Queries
Memon (Keys)
09.09
Memmo Aamoyyy (Kopp
Query Network
Conv
MLP
Graph Label
DOOG
The proposed architecture for hierarchical graph representation learning using the pro- posed memory layer. The query network projects the initial node features into a latent query space and each memory layer jointly coarsens the input queries and transforms them into new query space.
Figure 1:
ie., pooling, and transform their features, i.e., representation learning. As shown in Figure [ memory layer consists of arrays of memory keys, i.e. multi-head memory, and a convolutional layer. Assuming |h| memory heads, shared input query is compared against all the keys in each head resulting in / attention matrices which are then aggregated into a single attention matrix using the convolution layer.
In content addressable memory (Graves et al. 2014 Sukhbaatar et al. 2015 Weston et al. 2015), the task of attending to memory, i.e., addressing scheme, is formulated computing the similarity between memory keys to given query q. Specifically, the attention weight of key ký for query is defined as W/ so ta/amma (d(q, q, (;)) where d similarity measure, typically Euclidean distance or cosine similarity (Rae al.) 2016). The soft read operation memory is defined as a weighted average over the memory keys:
Ej Wykg..
In this work, we treat the input queries Q E Reexd the node representations of an input graph and treat the keys K()) -m++-xdd the cluster centroids of the queries. To satisfy this assumption, we impose clustering-friendyy distribution as the distance metric between keys and query. Following (Xie et al. [0166] Maaten & Hinton 2008), we use the Student's t-distribution as kernel measure the normalized similarity between query y and key kj as follows:
1 11g: - -k;]12/7) 7
Ci.
(1)
Σ 11 11g; kyy11 =fi
where Cip the normalized score between query yi and key ky, i.e., probability of assigning node to cluster j or attention score between query qs and memory key ky, and is the degree of freedom of the Student's t-distribution, i.e., temperature. To increase the capacity, we model the memory keys multi-head array. Applying shared input query against the memory keys produces tensor of cluster assignments _C. E R Xm1i1 where |/| denotes the number of heads.
To aggregate the heads into a single assignment matrix, we treat the heads and the assignments ma- trices depth, height, and width channels in standard convolution analogy and apply convolution operator over them. Because there is no spatial structure, we use x 1] convolution to aggregate the information across heads and therefore the convolution behaves weighted pooling that reduces the heads into a single matrix. The aggregated assignment matrix is computed as follows:
( softmax (ra
(1a(11) k=
[+]
c!")
E 1EX
(2)

---

Published as a conference paper at ICLR 2020
where o is a [1 x 1] convolutional operator parametrized by ø, is the concatenation operator, and C(1) the aggregated soft assignment matrix.
memory read generates a value matrix v() E R that represents the coarsened node repre- sentations in the same space as the input queries and is defined as the product of the soft assignment scores and the original queries as follows:
V C Q+)
(3)
The value matrix fed to single-layer beed-forward neural network to project the coarsened em- beddings from R2 xd: into R= xdi+1 as follows:
(a+1) =0 0a (vo))) E 1 xdx+1
(4)
where Q!++1) E 1 the output queries, W Rd is the network parameters, and is the non-linearity implemented using LeakyReLU.
Thanks to these parametrized transformations, a memory layer can jointly learn the node represen- tations and coarsens the graph end-to-end. The computed queries Q are the input queries to the subsequent memory layer ((++)) For graph classification tasks, one simply stack layers of memory up to the level where the imput graph is coarsened into a single node representing the global
graph representation and then feed it to fully-connected layer to predict the graph class follows:
softmax (MLP (MO (M(l-1) (...M())
(5)
where Q. (q) is the initial query representation generated generatee by the query network f, over graph g. We introduce two architectures based on the memory layer: GMN and MemGNN. These two architectures are different in the way that the query network is implemented. More specifically, GMN uses a ferwafd network for initializing the query: fg(g) = PFN+(g). whereas MemGNN implements the query network message passing GNN: f4(9) GNN
3.2 GMN ARCHITECTURE
GMN is a stack of memory layers on top of a query network fg(g) that generates the initial query representations without any message passing. Similar to set neural networks (Vinyals et al.) 2015) and transformers (Vaswani et al., 2017), nodes in GMN are treated as a permutation-invarte set of representations. The query network projects the initial node features into latent space that represents the initial query space.
Assume training set D 191. 191. .2..... gNv of N graphs where each graph is represented 9 (A, X, Y) and A E {0, 1|mx" denotes the adjacency matrix, X E gaxde is the initial node features, and Y E R is the graph label. Considering that GMN treats graph as set of order-invarian nodes and does not use message passing, and also considering that the memory layers do not rely on connectivity information, the topological information of each node should be somehow encoded into its initial representation. To define the topological embedding, we use an instantiation of general graph diffusion matrix S E IR=xn More specifically, we use random walk with restart (RWR) (Pan et al., 2004) to compute the topological embeddings and then sort them row-wise to force the embedding be irder-invariant. For further details please see section A..4] Inspired by transformers Vaswani et al. 2017). we then fuse the topological embeddings with the initial node features into the initial query representations using query network fy implemented as a two-layer feed-forward neural network:
Q (0) =o =o<[[[(So) x)]W))
Wi)
(6)
We use initial node representation and initial query representation interchangeably throughout the paper.
4

---

Published as a conference paper at ICLR 2020
where Wo E Ruxdn and W1 R2 the parameters of the query networks, || is the concatenation operator, and a is the non-linearity implemented using LeakyReLU.
3.2.1 PERMUTATION INVARIANCE
Considering the inherent permutation-invarian property of graph-structured data, model designed address graph classification tasks, should also enforce this property. This implies that the model should generate same outputs for isomorphic input graphs. We impose this on GMN architecture by sorting the topologiaal embedding row-wise as a pre- processing step.
Proposition 1. Given sorting function GMN (S), X) is permulationn-ivvaiinn.
Proof. Let P E {0, 13 (xxx)) be an arbitrary permutation matrix. For each node in graph G with adjacency matrix A, the corresponding node in graph Gp with permuted adjacency matrix PAP" has the permuted version of the topological embedding of the node in graph G. Sorting the embeddings cancels out the effect of permutation and makes the corresponding embeddings in graph G and Gp identical.
3.3 MEMGNN ARCHITECTURE
Unlike the GMN architecture, the query network in MemGNN relies on message passing to compute the initial query Qn as follows:
( = Go (A,X)
where query network Gp is an arbitrary parameterized message passing GNN (Gilmer et al. 2017
Li ei al. 2015 Kipf & Welling 2016 Velikovi et al.) 2018).
In our implementation, we use modified variant of GAT (Velikovi et al. 2018). Specifically, we introduce an extension to the original GAT model called edge-based GAT (e-GAT) and use it as the query network. Unlike GAT, e-GAT learns attention weights not only from the neighbor nodes but also from the input edge features. This is especially important for data containing edge information
(e.g., various bonds among atoms represented edges in molecule datasets). In e-GAT layer, attention score between two neighbor nodes is computed as follows:
W W W
|| W. W.h 1))
exp
Oii
(8)
Σ kEN:
exp
!1) [ W W
Wa (
W
where h. and h22 denote the representation of node i and the representation of the edge connecting node to its one-hop neighbor node j in layer respectively. W,, and W. are learnable node and edge weights and W the parameter a single-layer feed-forward network that computes the attention score. o is the non-linearity implemented using LeakyReLU.
3.4 TRAINING
We jointly train the model using two loss functions: supervised loss and an unsupervised clustering loss. The supervised loss denoted as Lrup is defined as cross-entropy loss and root mean square error (RMSE) for graph classification and regression tasks, respectively. The unsupervised clustering loss is inspired by deep clustering methods (Razavi et al. 2019 Xie et al.) 2016 Aljalbout et al.) [2018). embeddings the latent space by learning from high confidence assignments with the help of an auxiliary target distribution. The unsupervised
encourages the model to learn clustering-friente
loss is defined the Kullback-Leibler (KL) divergence between the soft assignments (()) and the auxiliary distribution p(!) follows:
5

---

Published as a conference paper at ICLR 2020
p!l
C KI.
(P((( (p())I(()
= =ΣE
log
For the target distributions p()) izes the loss contributions and improves the cluster purity while emphasizing on the samples with higher confidence. This distribution is defined as follows:
the distribution proposed in (Xie et al.
2016) which normal-
(C())
2
Σ
(10)
Σ
(
Σ
We define the total loss follows where L is the number of memory layers and À [0, 1] is scalar weight.
C
1 N
ALsup (1 A)
Σ
n KL
(11)
We initialize the model parameters, the keys, and the queries randomly and optimize them jointly with respect to L. using mini-batch stochastic gradient descent. stabilize the training, the gradients of Loup are back-propagated batch-wise while the gradients of L are applied epoch-wise by periodically switching X between and 1. Updating the centroids, i.e., memory keys, with the same frequency as the network parameters can destabilize the training. To address this, we optimize all model parameters and the queries in each batch with respect Loup and each epoch with
(!)) respect to KKT, Memory keys, on the other hand, are only updated at the end of each epoch by the
gradients of L This technique has also been applied in (Hassani & Haley, 2019 Caron et al, 2018) to avoid trivial solutions.
4 EXPERIMENTS
4.1 DATASETS
We use nine benchmarks including seven graph classification and two graph regression datasets to evaluate the proposed method. These datasets are commonly used in both graph kernel (Borgward)
& Kriegel 2005; Yanardag & Vishwanathan 2015) (Shervashidze al. 2009 Ying al. 2018
Shervashidze et al. 2011 Kriege et al. (2016) and GNN (Cangea et al. 2018: Ying et al. 2018 Lee
al. 2019 Gao & Ji. 2019) literature. The summary of the datasets is follows where the first two benchmarks are regression tasks and the rest are classification tasks.
ESOL (Delaney (2004) contains water solubility data for compounds. Lipophilicity (Gaulton et al. 2016) contains experimental results of octanol/water distribution of compounds.
Bace (Subramanian et al. 2016 provides quantitative binding results for set of inhibitors of human
8-secretase (BACE-1)
DD (Dobson & Doig. 2003) is used to distinguish enzyme structures from non-enzymes.
Enzymes (Schomburg et al. (004) for predicting functional classes of enzymes.
Proteins (Dobson & Doig 2003) used to predict the protein function from structure.
Collab (Yanardag Vishwanathan 2015) for predicting the field of a researcher given one's ego-
collaboration graph.
Reddit-Binary (Yanardag & Vishwanathan 2015) is for predicting the type of community given a graph of online discussion threads.
Tox21 (Challenge) 2014) is for predicting toxicity on 12 different targets.
For more information about the detailed statistics of the datasets refer to Appendix A.2

---

Published as a conference paper at ICLR 2020
Table 1: Mean validation accuracy over 10-folds.
Method
Dataset
Enzymes
Proteins
DD
Collab
Reddit-B
Graphlet (Shervashidze et al., 2009
41.03
72.91
64.66
64.66
78.04
Kernel
ShortestPath (Borgwardt & K Kriegel
2005
42.32
76.43
78.86
59.10
64.11
WL (Shervashidze et al. [1,]]0011
53.43
73.76
74.02
78.61
68.20
WT. Optimal (Kriege al. 2016)
60.13
75.26
79.04
80.74
89.30
PatchySan (Niepert et al. 2016)
75.00
76.27
72.60
86.30
GraphSage (Hamilton et al.
2017
54.25
70.48
75.42
68.25
ECC (Simonovsky & Komodakis
2017)
53.50
72.65
74.10
67.79
Set2Set Vinyals el al.
2015)
60.15
74.29
78.12
71.75
SortPool (Morris el al. 2019)
57.12
75.54
79.37
73.76
GNN
DiffPool (Ying et al. 2018)
60.53
76.25
80.64
75.48
85.95
CliquePool (Luzhmica et al
2019
60.71
72.59
77.33
74.50
Sparse HGC (Cangea et al.
2018
64.17
75.46
78.59
75.46
79.20
TopKPool (Gao & 2019
77.68
82.43
77.56
74.70
SAGPool (Lee et al. 2019
71.86
76.45
73.90
GMN (ours)
78.66
82.25
84.40
80.18
95.28
MemGNN (ours)
75.50
81.35
82.92
77.0
85.55
4.2 GRAPH CLASSIFICATION RESULTS
To evaluate the performance of our models on DD, Enzymes, Proteins, Collab, and Reddit-Binn Binaary datasets, follow the experimental protocol in (Ying et al.] 2018)) and perform 10-fold cross- validation and report the mean accuracy over all folds. We also report the performance of four graph kernel methods including Graphlet (Shervashidze et al. 2009), shortest path (Borgwardt & Kriegel, 2005). Ceisfeiler-Lehmee (WL) (Shervashidze et al.) 2011). and WL Optimal Assignment (Kriege al. 2016), and ten GNN models.
The results shown in Table suggest that: (i) our models achieve state-of-the-ar results w.r.t. GNN models and significantly improve the performance Enzymes, Proteins, DD, Collab, and Reddit- Binary datasets by absolute margins of 14.49% 6.0%, 3.76%, 2.62%, and 8.98% accuracy, respec- tively, (ii) our models outperform graph kernels on all datasets except Collab where our models competitive with the best kernel, i.e., absolute margin of 0.56%, (iii) both proposed models achieve better performance or are competitive compared to the baseline GNNs, (iv) GMN achieves better re- sults compared to MemGNN which suggests that replacing local adjacency information with global topological embeddings provides the model with more useful information, and (v) On Collab, our models are outperformed by variant of DifiPpool (i.e., diffpool-det) (Ying et al., 2018) and WI. Optimal Assignment (Kriege et al. 2016). The former a GNN augmented with deterministic clus- tering algorithm whereas the latter is graph kernel method. We speculate this is because of the high edge-to-node ratio in this dataset and the augmentations used these two methods help them with extracting near-optimal cliques.
To evaluate the performance on the BACE and Tox21 datasets, we follow the evaluation protocol in (Wu et al.) 2018) and report the under the curve receiver operating characteristics (AUC-ROC) measure. Considering that the BACE and Tox21 datasets contain initial edge features, we train the MemGNN model and compare its performance the baselines reported in (Wu al.] 2018). The results shown in Table suggest that our model achieves state-of-the-an results by absolute margin of 4.0 AUC-ROC the BACE benchmark and competitive with the state-of-the-art GCN model on the Tox21 dataset, i.e., absolute margin of 0.001.
4.3 GRAPH REGRESSION RESULTS
For the ESOL and Lipophilicity datasets, we follow the evaluation protocol in (Wu et al. 2018) and report their RMSEs. Considering that these datasets contain initial edge features (refer to Appendix A.2 for further details), train the MemGNN model and compare the results the baseline mod-
In diffpool-det assignment matrices generated using deterministic graph clustering algorithm.

---

Published as a conference paper at ICLR 2020
Table 2: AUC-ROC on BACE and Tox21 datasets.
Method
Dataset
BACE
Tox21
validation
test
validation
test
Logistic Regression
0.719 = 0.003
0.781 = 0.010
0.772 = 0.011
0.794 0.015
KernelS
0.739 + 0.000
0,862 + 0.000
0.818 + 0.010
0.822 + 0.006
XGBoost
0.756 + 0.000
0.850 + 0.008
0.775 + 0.018
0.794 0.014
Random Forest
0.728 + 0.004
0.867 + 0.008
0.763 + 0.002
0.769 + 0.015
IRV
0.715 + 0.001
0.838 + 0.000
0.807 = 0.006
0.799 0.006
Multitask
0.696 0.037
0.824 0.0006
0.795 0.017
0.803 1003=:00 0.012
Bypass
0.745 + 0.017
0.829 0.006
0.800 + 0.008
0.810 0.013
GCN
0.627 = 0.015
0.783 + 0.014
0.825 0.013
0.829 = 0.006
Weave
0.638 + 0.014
0.806 + 0.002
0.828 + 0.008
0.820 + 0.010
MemGNN (ours)
0.859 = 0.000
0.907 = 0.000
0.862 = 0.009
0.828 = 0.004
Table 3: RMSE on ESOL and Lipophilicity datasets.
Method
Dataset
ESOL
Lipophilicity
validation
test
validation
test
Multitask
1.17 0.13
1.22 0.19
0.852 0.048
0.859 = 0.013
Random Forest
1.16 0.15
1.07 + 0.19
0.835 + 0.036
0.876 + 0.040
XGBoost
1.05
0.99
0.783
0.021
0.799 + 0.054
GCN
0.15
0.97
0,678 + 0.040
= #
MPNN
0.02
0.00
0.757 0.030
0.715 = 0.035
KRR
1.65 0.19
0.06
0.889 = 0.009
0.899 = 0.043
DAG
0.74 0.04
0.02
0.857 0.050
0.835 = 0.039
Weave
0.57 0.04
0.61 +
0.734 + 0.011
0.715 + 0.035
MemGNN (ours)
0.53 + 0.03
0.54 0.01
0.555 + 0.039
0.556 = 0.023
els reported in (Wu et al., 2018) including graph based methods such as GCN, MPNN, directed acyclic graph (DAG) model, and Weave as well as other conventional methods such kernel ridge regression (KRR) and influence relevance voting (IRV). Results shown in Table 3 suggest that our MemGNN model achieves state-of-the-art results by absolute margin of 0.07 and 0.1 RMSE ESOL and Lipophilicity benchmarks, respectively. For further details on regression datasets and baselines please refer to (Wu al. 2018))
4.4
ABLATION STUDY
4.4.1 EFFECT OF EDGE FEATURES
To investigate the effect of the proposed e-GAT model, we train the MemGNN model using both GAT and e-GAT layers as the query network. Considering that the ESOL,, Lipophilicity, and BACE datasets contain edge features, compared to edges, we set the node and edge feature dimensions to and 4, respectively. The performance of the two layers on the ESOL dataset shown in Appendi (v.ediiXA.]] suggesting that e-GAT achieves better results on the validation set in each epoch compared to the standard GAT model. We observed the same effect Lipophilicity and BACE datasets.
them the benchmarks. Since nodes have richer features
4.4.2 EFFECT OF TOPOLOGICAL EMBEDDING
To investigate the effect of topological embeddings on the GMN model, we evaluated three initial topological features including adjacency matrix, normalized adjacency matrix, and RWR. For fur- ther details on RWR, see section A.4] The results suggest that using the RWR as the initial positional embedding achieves the best performance. For instance, 10-fold cross validation accuracy of GMN model trained on Enzymes dataset with adjacency matrix, normalized adjacency matrix, and RWR 78.66%, 77.16 77.16% 6%, and 77.33%, respectively. Furthermore, sorting the topological embeddings
8

---

Published as a conference paper at ICLR 2020
(a)
(b)
Figure 2: Visualization of the learned clusters of two molecule instances from (a) ESOL and (b) Lipophilicity datasets. The visualizations show that the learned clusters correspond to known chem- ical groups. Note that node without label represents a carbon atom. For more visualizations and discussion see section|..
to guarantee invariance to permutations improves the performance. For example, it increases the accuracy on the DD dataset from 82.24% to 84.40%.
4.4.3 DOWN-SAMPLING NEIGHBORS WITH RANDOM WALKS
We investigate two methods to down-sample the neighbors in dense datasets such as Collab (i.e., average of 66 neighbors per node) to enhance the memory and computation. The first method ran- domly selects 10% of the edges whereas the second method ranks the neighbors based on their RWR scores with respect the center node and then keeps the top 10% of the edges. We trained the MemGNN model on Collab using both sampling methods which resulted in 73.9% and 73.1% % 10-fold cross validation accuracy for random and RWR-based sampling methods respectively, sug- gesting that random sampling performs slightly better than RWR based sampling.
4.4.4 EFFECT OF NUMBER OF KEYS AND HEADS
We speculate that although keys represent the clusters, the number of keys is not necessarily pro- portional to the number of the nodes in the input graphs. In fact, datasets with smaller graphs might have more meaningful clusters to capture. For example, molecules are comprised of many functional groups and yet the average number of nodes the ESOL dataset 13.3. Moreover, our experiments show that for Enzymes with average number of 32.69 nodes, the best performance is achieved with 10 keys whereas for the ESOL dataset 64 keys results in the best performance. 160 keys result in RMSE of 0.56, 0.52, and 0.54, respectively. We also observed that with fixed number parameters, increasing the number of memory heads improves the performance. For in- stance, when the model is trained on ESOL with 160 keys and head, it achieves RMSE of 0.54, whereas when trained with 32 keys and 5 heads, the same model achieves RMSE of 0.53.
ESOL 8, 64, and
4.4.5 WHAT DO THE KEYS REPRESENT?
Intuitively, the memory keys represent the cluster centroids and enhance the model performance by capturing meaningful structures. To investigate this, we used the learned keys to interpret the knowledge learned by models through visualizations. Figure 2 visualizes the learned clusters over atoms (i.e., atoms with the same color are within the same cluster) indicating that the clusters mainly consist of meaningful chemical substructures such as a carbon chain and a Hydroxyl group (OH) (i.e., Figure 22),, as well as a Carboxyl group (COOH) and a benzene ring (i.e., Figure 2b). From chemical perspective, Hydroxyl and Carboxyl groups, and carbon chains have significant impact on the solubility of the molecule in water or lipid. This confirms that the network has learned chemical features that essential for determining the molecule solubility. It is noteworthy that tried initializing the memory keys using K-Means algorithm over the initial node representations to wart-start them but did not observe any significant improvement over the randomly selected keys.

---

Published as a conference paper at ICLR 2020
5 CONCLUSION
We proposed an efficient memory layer and two models for hierarchical graph representation learn- ing. We evaluated the proposed models nine graph classification and regression tasks and achieved state-of-the-art results on eight of them. We also experimentally showed that the learned representations capture well-known chemical features of the molecules. Furthermore, showed that concatenating node features with topological embeddings and passing them through few memory layers achieves notable results without using message passing. Moreover, we showed that defining the topological embeddings using graph diffusion achieves best performance. Finally, we showed that although connectivity information not explicitly imposed on the model, the mem- ory layer can process node representations and properly cluster and aggregate the learned represen- tations.
we discussed that a graph kernel and a GNN augmented with determin- istic clustering achieve better performance compared to our models on the Collab dataset. Analyzing samples in this dataset suggests that in graphs with dense communities, such as cliques, our model struggles to properly detect the dense sub-graphs.
Limitations: In section 4.
We plan to extend our models to also perform node classification by attending to the node representations and centroids of the clusters from different layers of hierarchy that the nodes belongs to. Moreover, alized PageRank and heat kernel, to initialize the topological embeddings. We are also planing to investigate the representation learning capabilities of the proposed models in self-suue setting.
Future Directions:
planning evaluate other graph diffusion models, e.g., person-
REFERENCES
Elie Aljalbout, Vladimir Golkov, Yawar Siddiqui, Maximilian Strobel, and Daniel Cremers. Clus- tering with deep learning: Taxonomy and new methods. arXiv preprint arXiv: 881.00668, 2018.
Mikhail Belkin and Partha Niyogi Laplacian eigenmaps and spectral techniques for embedding and clustering. In Advances in Neural Information Processing Systems, pp. 585-591. 2002.
Filippo M. Bianchi, Daniele Grattarola, and Cesare Alippi. Mincut pooling in graph neural networks. arXiv preprint arXiv: 1907.0048,, 2019.
Karsten M Borgwardt and Hans-Peter Kriegel. Shortest-path kernels on graphs. In International Conference on Data Mining, pp. 8 8-pp. IEEE, 2005.
Joan Bruna, Wojciech Zaremba, Arthur Szlam, and Yann LeCun. Spectral networks and locally connected networks on graphs. In International Conference on Learning Representation, 2014.
Cätälina Cangea, Petar Veličković, Nikola Jovanović, Thomas Kipf, and Pietro Liò. Towards sparse hierarchical graph classifiers. In Advances in Neural Information Processing Systems, Workshop on Relational Representation Learning, 2018
Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze. Deep clustering for unsu- pervised learning of visual features. In European Conference on Computer Vision, pp. 132-149, 2018.
Tox21 Data Challenge. Tox21 data challenge 2014, 2014.
Hanjun Dai, Bo Dai, and Le Song. Discriminative embeddings of latent variable models for struc- tured data. In International Conference on Machine Learning, pp. 2702-2711, 2016.
Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. In Advances Neural Information Processing Systems, pp. 3844 3852, 2016.
John S Delaney. Esol: estimating aqueous solubility directly from molecular structure. Journal of Chemical Information and Computer Sciences, 44(3):1000-1000 2004.
Inderjit S Dhillon, Yuqiang Guan, and Brian Kulis. Weighted graph cuts without eigenvectors multilevel approach. IEEE Transactions on Pattern Analysis and Machine Intelligence, 29(11); 1954-1957, 2007.
10

---

Published as a conference paper at ICLR 2020
Frederik Diehl. WwXi:1190..109000 2019.
Edge contraction pooling for graph neural networks.
arXiv preprint
Paul D Dobson and Andrew Doig. Distinguishing enzyme structures from non-enzymes without alignments. Journal of Molecular Biology, 330(4):771-783, 2003.
David K Duvenaud, Dougal Jorge Iparraguirre, Rafael Bombarell, Timothy Hirzel, Alán Aspuru-Guzik, and Ryan P Adams. Convolutional networks on graphs for learning molecular fingerprints. In Advances in Neural Information Processing Systems, pp. 2224-2232, 2015.
Hongyang Gao and Shuiwang Ji. Graph u-nets. In International Conference Machine Learning. pp. 2083-2092, 2019.
Anna Gaulton, Anne Hersey, Michal Nowotka, A Patrícia Bento, Jon Chambers, David Mendez, Prudence Mutowo, Francis Atkinson, Louisa J Bellis, Elena Cibrián-Uhalte, al. The chembl database in 2017. Nucleic Acids Research,
2016.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning, pp- 1263-1272, 2017.
Alex Graves, Greg Wayne, and Ivo Danihelka. arXiv:1410.5401, 2014.
Neural turing machines.
arXiv preprint
Aditya Grover and Jure Leskovec. Node2vee: Scalable feature learning for networks. Interna- tional Conference on Knowledge Discovery and Data Mining, pp. 855-864, 2016.
Will Hamilton, Zhitao Ying. and Jure Leskovec. Inductive representation learning on large graphs. In Advances in Neural Information Processing Systems, pp. 1024-1034, 2017.
Kaveh Hassani and Mike Haley. Unsupervised multi-task feature learning on point clouds. In International Conference Computer Vision, pp. 8160-8871, 8171, 2019.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pp. 448-456, 2015.
Diederik P Kingma and Jimmy Lei Ba. Adam: Amethod for stochastic optimization. In Interna- tional Conference on Learning Representation, 2014.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net-
works. In International Conference on Learning Representations, 2016.
Nils M Kriege, Pierre-Louis Giscard, and Richard Wilson. On valid optimal assignment kernels and applications to graph classification. In Advances in Neural Information Processing Systems, pp- 1623 1631, 2016.
Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. arXiv preprint arXiv: 2019.
Junhyun Lee, Inyeop Lee, and Jaewoo Kang. Self-attention graph pooling. In International Confer- ence on Machine Learning, pp. 3734-3743, 2019.
Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper insights into graph convolutional networks for semi-supervised learning. In AAAI Conference Artificial Intelligence, pp. 3538-3545, 201 8.
Yujia Li, Daniel Tarlow, Mare Brockschmidt, and Richard Zemel. Gated graph sequence neural networks. In International Conference on Learning Representations, 2015.
Enxhell Luzhnica, Ben Day, and Pietro Lio. Clique pooling for graph classification. In International Conference on Learning Representations, Workshop Representation Learning Graphs and Manifolds, 2019.
Jianxin Ma, Peng Cui, Kun Kuang, Xin Wang, and Wenwu Zhu. Disentangled graph convolutional networks. In International Conference on Machine Learning, pp- 4212 4221, 2019.

---

Published as a conference paper at ICLR 2020
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of Machine Learning Research, 9:2579-2605, 2008.
Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In AAAI Conference on Arlificial Intelligence, volume 33, pp. 4602-4609, 2019.
Mathias Niepert, Mohamed Ahmed, and Konstantin Kutzkov. Learning convolutional neural net- works for graphs. In International Conference on Machine Learning, pp. 2014-2023, 2016.
Mingdong Ou, Peng Cui, Jian Pei, Ziwei Zhang, and Wenwu Zhu. Asymmetric transitivity preserv- ing graph embedding. In International Conference on Knowledge Discovery and Data Mining, pp. 1105-1114, 14, 2016.
Jia-Yu Yu Pan, Hyung-Jeon Yang. Christos Faloutsos, and Pinar Duygulu. Automatic multimedia cross-modal correlation discovery. In International Conference on Knowledge Discovery and Data Mining. pp. 653-655 653, 2004.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Vetioo, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Bryan Perozzi, Rami AI-Rfou, and Steven Skiena. Deepwalk: Online learning of social represen- tations. In International Conference on Knowledge Discovery and Data Mining, pp. 701-710, 2014.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adria Puigdomenech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In International Conference on Machine Learning, pp. 2827-2836, 2017.
Jack Rac, Jonathan Hunt, Ivo Danihelka, Timothy Harley, Andrew W Senior, Gregory Wayne, Alex Graves, and Timothy Lillicrap. Scaling memory-augman neural networks with sparse reads and writes. In Advances Neural Information Processing Systems, pp. 3621-3629, 2016.
Ali Razavi, Aaron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. arXiv preprint arXiv:/9066 00446, 2019.
Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta- learning with memory-augmente neural networks. In International Conference Machine Learning, pp. 1842-1850, 2016.
Michael Schlichtkrull, Thomas N Kipf, Peter Bloem, Rianne Van Den Berg, Ivan Titov, and Max Welling, Modeling relational data with graph convolutional networks. In European Semantic Web Conference, pp. 593-607. Springer, 2018.
Ida Schomburg, Antje Chang. Christian Ebeling, Marion Gremse, Christian Heldt, Gregor Huhn, and Dietmar Schomburg. Brenda, the enzyme database: updates and major new developments. Nucleic Acids Research, 32:D431 2004.
Nino Shervashidze, SVN Vishwanathan, Tobias Petri, Kurt Mehlhorn, and Karsten Borgwardt. Ef- ficient graphlet kernels for large graph comparison. In Artificial Intelligence and Statistics, pp- 488-495, 2009.
Nino Shervashidze, Pascal Schweitzer, Erik Jan van Leeuwen, Kurt Mehlhorn, and Karsten M Borg- wardt. Weisfeiler-Iehman graph kernels. Journal of Machine Learning Research, 12(Sep):2539- 2561,2111.
Martin Simonovsky and Nikos Komodakis. Dynamic edge-conditioned filters in convolutional neu- ral networks on graphs. In International Conference on Computer Vision and Pattern Recognition, pp. 3693-3702, 2017.
Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. In Advances in Neural Information Processing Systems, pp. 4077-4087, 2017.
12

---

Published as a conference paper at ICLR 2020
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, 1lya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15:1929-1955 1958, 2014.
Govindan Subramanian, Bharath Ramsundar, Vijay Pande, and Rajiah Aldrin Denny. Computational modeling of ß-secretase 1 (bace-1) inhibitors using ligand based approaches. Journal of Chemical Information and Modeling, 56(10): (100):193-111)) 1949, 2016.
Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. End-to-end memory networks. In Advances in Neural Information Processing Systems, pp. 2440-2448, 2015.
Hanghang Tong, Christos Faloutsos, and Jia- Yu Pan. Fast random walk with restart and its applica- tions. In International Conference Data Mining, pp. 613-622, IEEE, 2006.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances Neural Infor- mation Processing Systems, pp. 5998-6008 2017.
Petar Velikovi, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Li, and Yoshua Ben- gio. Graph attention networks. In International Conference on Learning Representations, 2018.
Oriol Vinyals, Samy Bengio, and Manjunath Kudlur. Order matters: Sequence to sequence for sets.
In International Conference on Learning Representations, 2015.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. Matching networks for one shot learning. Advances in Neural Information Processing Systems, pp. 3630-3638, 2016.
Salvatore Vivona and Kaveh Hassani. Relational graph representation learning for open-domain question answering, arXiv preprint arXiv:1910.08244 2019.
Chu Wang, Babak Samari, and Kaleem Siddiqi. Local spectral graph convolution for point set feature learning. In European Conference on Computer Vision, pp. 52-66, 2018.
Tingwu Wang, Yuhao Zhou, Sanja Fidler, and Jimmy Ba. Neural graph evolution: Automatic robot design. In International Conference Learning Representations, 2019.
Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. In International Conference on Learning Representation, 2015.
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: benchmark for molecular machine learn- ing. Chemical Science, 92):513-530, 2018
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and Philip S Yu. A comprehensive survey graph neural networks. arXiv preprint arXiv: 1901100596, 2019.
Junyuan Xie, Ross Girshick, and Ali Farhadi. Unsupervised deep embedding for clustering analysis. In International Conference on Machine Learning, pp. 478 487, 2016.
Keyulu Xu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe. Ken-ichi Kawarabayashi, and Stefanie Jegelka. Representation learning on graphs with jumping knowledge networks. In International Conference on Machine Learning, pp. 5453-546 2018.
Pinar Yanardag and S.V.N. Vishwanathan. Deep graph kernels. International Conference on Knowledge Discovery and Data Mining, pp. 1365-1374, 2015.
Zhitao Ying, Jiaxuan You, Christopher Morris, Xiang Ren, Will Hamilton, and Jure Leskovec. Hi- erarchical graph representation learning with differentiable pooling. In Advances in Neural Infor- mation Processing Systems, pp. 4800-488 4810, 2018
Ziwei Zhang, Peng Cui, and Wenwu Zhu. Deep learning on graphs: A survey, arXiv preprint arXiv:1812.04220 2018.
Jie Zhou, Ganqu Cui, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, and Maosong Sun. Graph neural networks: A review of methods and applications. arXiv preprint 0844401 2018.
13

---

Published as a conference paper at ICLR 2020
A APPENDIX
A.1 IMPLEMENTATION DETAILS
We implemented the model with PyTorch (Paszke al,, 2017) and optimized it using Adam (Kingma & Ba 2014) optimizer. We trained the model for a maximum number of 2000 epochs and decayed the learning rate every 500 epochs by 0.5. The model uses batch-normalization (loffe & Szegedy 2015), skip-connections, LeakyRelu activation functions, and dropout (Srivastava al., 2014) for regularization. We also set the temperature Students I-distribution to 1.0 and restart probability in RWR to 0.1. We decided the hidden dimension and the number of model parame- ters using random eyper-parameter search strategy. The best performing rrper-parameters for the datasets are shown in Table 4
A.2 DATASET STATISTICS
Table 5
summarizes the statistics of the datasets used for graph classification and regression tasks.
A.3 EFFECT OF E-GAT
section 4.4.1 we introduced e-GAT. Figures 3a and 3b illustrate the RMSE and R² score on validation set of the ESOL dataset achieved by a MemGNN model using both GAT and e-GAT the query network, respectively. As shown, e-GAT performs better compared to GAT on both metrics.
A.4 RANDOM WALK WITH RESTART
Suppose an agent randomly traverses a graph starting from node and iteratively walks towards its neighbors with a probability proportional to the edge weight that connects them. The agent also can randomly restart the traverse with probability p. Eventually, the agent will stop at node j with probability called relevance score of node j with respect to node i (Tong et al. 2006). The relevance score of node with every other node of the graph is defined as follows:
T PAC, + (1- = (1- -À)-
(12)
where f is the RWR score corresponding to node i, p is the restart probability, Á is the normalized adjacency matrix, and e is one-hot vector representation of node
Note that the restart probability defines how far the agent can walk from the source node and there- fore E represents the trad-off between local and global information around node i.
A.5 LEARNED CLUSTERS
Figure shows how unsupervised loss helps the model to push the nodes into distinct clusters. Figures 4a and 4c illustrates clusters with unsupervised loss and Figures 4b and 4d show com- puted clusters without unsupervised loss. The visualizations suggest that unsupervised loss helps the model to avoid trivial solutions by collapsing the latent node representations into meaningless meaningful chemical groups extracted by MemGNN. Figures 5b and and ade from L.PO and figure 5a and 5c are from ESOL dataset respectively.
clusters. Also,
14

---

Published as a conference paper at ICLR 2020
16-
GAT
13
GAT
12
RMSE
85-
GAT GAT
85
200
900
6.0
6:0
1000
:30e
180e
200
600
100
1000
1000
1000
Step
Step
(a)
(b)
Figure 3: Validation (a) R² score, and (b) RMSE achieved by MemGNN model on ESOL with GAT and e-GAT based query networks.
(b)
(c)
(d)
Figure 4: Figures (b) and (d) show computed clusters without using unsupervised clustering loss, whereas Figures (a) and (c) show the clusters learned using the unsupervised clustering loss. The visualizations suggest that the unsupervised loss helps the model in learning distinct and meaningful clusters.
15

---

Published as a conference paper at ICLR 2020
Table 4: Hyper-parametce selected for the models.
Dataset
#Keys
#Heads
#Layers
Hidden Dimension
Batch
Enzymes
[10.1]
5
100
20
Proteins
[10,1]
5
2
80
20
DD
[16, 8, 1]
5
120
64
Collab
[32, 8, 1]
5
3
100
64
Reddit-B
[32,1]
I
2
16
32
ESOL
[64,
5
2
16
32
Lipophilicity
[32,1]
5
2
16
32
BACE
[32.1]
5
8
32
Table 5: Summary of statistics of the benchmark dataset.
Name
Task
Graphs
Classes
Avg. Nodes
Avg. Edges
Node Attr.
Edge Attr.
Enzymes
classification
600
6
32.63
62.14
18
0
Proteins
classification
1113
2
39.06
72.82
29
0
DD
classification
1178
2
284.32
715.66
0
0
Collab
classification
5000
3
74.49
2475.78
Reddit-B
classification
2000
2
429.63
497.75
0
0
Bace
classification
1513
2
34.09
36.86
32
7
Tox21
classification
8014
2
17.87
18.50
32
7
ESOL
regression
1144
13.29
13.68
Lipophilicity
regression
4200
27.04
29.50
32
7
(a)
(b)
(c)
(d)
Figure 5: Clusters learned by McMGNN for ESOL and LIPO dataset. Chemical groups like OH (hydroxyl group), CCI3, COOH (carboxyl group), CO (ketone group) as well as benzene rings have been recognized during the learning procedure. These chemical groups are highly active and have great impact on the solubility of molecules.
16


Reviews:
[{"review_id": "rJgM2hZeqB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes to use memory network (self-attention) to summarize information from graphs. The proposed approach \"soft clusters\" the node embeddings into coarser representation and eventually fed to a final MLP for classification. The empirical results on some standard datasets show promising gains of the algorithm.\n\nThe proposed approach stacks a few layers of self-attention on top of either some features of the nodes (including projected edge connections, parametrized by W_0) or the node embeddings of some form of graph neural network. And this stacking seems to be simple combination of existing approaches, without fully integrating them. In fact, the training process is also separated: \"task-specific loss are back-propagated batch-wise while the gradients of the unsupervised loss are applied\nepoch-wise\". It makes me wonder whether we can just separate it into two stages, i.e., first learn a node embedding using graph neural network, then learn this self attention transformation. Due to the above issues, I feel the novelty of the approach is limited and incremental.\n\nAnother issue with the paper is that the notations seem to be messed up and some concepts are not explained clearly. For example, the C_{i,j} soft assignment matrix is normalized row-wise, then Eqn (3) seems very suspicious, because it averages the queries using weights along the other direction, thus not normalized. Also, the dimension of the MLP weights do not align well with inputs, for instance in Eqn (4), it should be written as V^(l) W.\n\nThere are more questions that are not clearly specified in the current manuscript. For example, where does the keys K come from? From the text description, it seems to be cluster results, and do you do the clustering on every gradient update? Or are they learned from scratch? The distribution P defined in Eqn (11) also seems to be difficult to optimize since it depends on C_{i,j} and is connected to different entries. Is simple SGD sufficient to optimize over P? Moreover, in the experiment section, it is unknown how many layers of self-attention is applied and what are the important parameters. For better comparison, the experiment section should include some estimate of parameter size as well.\n\nThe experiment results seem interesting since the approach indeed achieves good performance across many datasets. Also the visualized keys are interesting as well because it captures some meaningful patterns from the data.\n\nThere is some related work that you should cite:\nHanjun Dai, Bo Dai and Le Song. Discriminateive Embeddings of Latent Variable Models for Structured Data. International Conference on Machine Learning (ICML) 2016.\n\n\n============================================================\nBased on the authors' reply and other reviews, I have changed my rating to \"Weak accept\".\n\nThe authors' reply has clarified the important detail of how the key matrix is learned and now it is clear that the algorithm is not just two-stage separate learning. In light of this clarification, I think the proposed algorithm is novel enough and the jointly training mechanism is also beneficial for the state-of-the-arts results reported in the experiments.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJeuzS4ucB", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a memory layer for Graph Neural Networks (GNNs) and two deep models for hierarchical graph representation learning. The proposed memory layer models the memory as a multi-head array of keys with a soft clustering mechanism and applies a convolution operator over the heads. The proposed models are experimentally evaluated on seven graph classification and regression tasks. \n\nGenerally, the paper is technically justified. The proposed technique is well motivated and properly presented. A novel clustering-convolution mechanism is proposed  for memory augmentation and graph pooling. However, there are still some rebuttal requests. 1- Some details are insufficient. For the multi-head mechanism, it is not stated clearly whether for each head an independent query is computed or a shared query is used for all heads.  \n2- Additionally, a related work published in NIPS 2016 should be cited and discussed. \nJack Rae et al. Scaling memory-augmented neural networks with sparse reads and writes.  \n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "B1xbJPHpYH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper presents \"memory layer\" to simultaneously do graph representation learning and pooling in a hierarchical way. It shares the same spirit with the previous models (DiffPool and Mincut pooling) which cluster nodes and learn representation of the coarsened graph. In DiffPool, Graph convolutional Neural Networks (GCNs) with 2-6 iterations of “Message passing” is used to learn node embedding, followed by graph pooling. By contrast, the proposed model circumvent the inefficiency of using message passing by using their proposed memory layer.\n\nPros:\n* The paper is generally written well, but some important details are missing.\n* Clear visualization of the results.\n* Useful ablation study.\n\nCons:\n* Missing information, which can be critical for the success of the model:\n(a) the estimation of the keys, and \n(b) how does the convolutional layer in Eq (2) work, given that the input for it is concatenation of matrices, which has no spatial structure?\n* Experiments on graph classification lack diversity, where CoLLAB is the only non-chemical dataset in the experiment. \n* The paper argues that interactive message passing is not efficient. But do you have any explanation on why MemGNN with message passing in initial embedding learning performs better than GMN without message passing in D&D dataset?\n* I have some reservation for calling something \"memory\", which is meant to store information for later processing. For this work, the network is a feed-forward architecture for processing graphs, where the middle layers (the queries) are matrices, which can be studied on their own right (e.g., see [1]).\n\nAt this point, the ideas for graph representation are plentiful, but there have not been a coherent story on how and why new architectures should work better than previous ones. This paper can be made stronger by offering insights along this line.\n\n[1] Do, K., Tran, T., & Venkatesh, S. (2017). Learning deep matrix representations. arXiv preprint arXiv:1703.01454.\n\n\n", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "B1xaBdyN9S", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors introduce a method for adding memory layers to a graph neural network which can be used for representation learning and pooling. Two variants, the MemGNN and GMN are proposed which use the memory layer. The authors evaluate their models over 7 datasets that cover classification and regression tasks. They obtain SOTA on 6/7 datasets; perform ablation analysis and introspect the clusters for chemical significance.\n\nOverall this paper is well written and easy to read. The motivation, equations and illustrations are clear and helpful. The model is technically novel, building up from existing approaches in a progressive way. Given the generality of the approach, the impact is also likely to be high.\n\nIn order to bolster their results the authors may run their approach on a few other datasets in Wu et. al. 2018. \n\nMinor issues:\n - Provide error bars for the tables\n - Sec 4.2 typo : “datastes” \n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Four reviewers have assessed this paper and they have scored it as 6/6/6/6 after rebuttal. Nonetheless, the reviewers have raised a number of criticisms and the authors are encouraged to resolve them for the camera-ready submission.
Decision: Accept (Poster)