Title: Graph Bernoulli Pooling
Venue: NeurIPS 2023
Accepted: no
Keywords: graph classification

variational inference

Bernoulli sampling

mutual information
Abstract: Graph pooling is crucial for enlarging receptive field and reducing computational cost in deep graph representation learning. In this work, we propose a simple but effective non-deterministic graph pooling method, called graph Bernoulli pooling (BernPool), to facilitate graph feature learning.
In contrast to most graph pooling methods with deterministic modes, we design a probabilistic Bernoulli sampling to reach an expected sampling rate through deducing a variational bound as the constraint. To further mine more useful info, a learnable reference set is introduced to encode nodes into a latent expressive probability space. Hereby the resultant Bernoulli sampling would endeavor to capture salient substructures of the graph while possessing much diversity on sampled nodes due to its non-deterministic manner. Considering the complementarity of node dropping and node clustering, further, we propose a hybrid graph pooling paradigm to combine a compact subgraph (via dropping) and a coarsening graph (via clustering), in order to retain both representative substructures and input graph info. Extensive experiments on multiple public graph classification datasets demonstrate that our BernPool is superior to various graph pooling methods, and achieves state-of-the-art performance. The code is publicly available in an anonymous format at \href{https://anonymous.4open.science/r/BernPool}{https:/github/BernPool}.

Abstract

Graph pooling is crucial for enlarging receptive field and reducing computational cost in deep graph representation learning. In this work, we propose a simple but effective non-deterministic graph pooling method, called graph Bernoulli pooling (BernPool), to facilitate graph feature learning. In contrast to most graph pooling methods with deterministic modes, we design a probabilistic Bernoulli sampling to reach an expected sampling rate through deducing a variational bound as the constraint. To further mine more useful info, a learnable reference set is introduced to encode nodes into a latent expressive probability space. Hereby the resultant Bernoulli sampling would endeavor to capture salient substructures of the graph while possessing much diversity on sampled nodes due to its nondeterministic manner. Considering the complementarity of node dropping and node clustering, further, we propose a hybrid graph pooling paradigm to combine a compact subgraph (via dropping) and a coarsening graph (via clustering), in order to retain both representative substructures and input graph info. Extensive experiments on multiple public graph classification datasets demonstrate that our BernPool is superior to various graph pooling methods, and achieves state-ofthe-art performance. The code is publicly available in an anonymous format at https:/github/BernPool.

1 Introduction

Graph Neural Networks (GNNs) [13, 36] have been widely used to learn expressive representation from ubiquitous graph-structured data such as social networks [31], chemical molecules [15] and biological networks [24]. To improve representation ability, multiple GNN variants, e.g., graph convolutional networks (GCNs) [16] and graph attention networks (GATs) [25], have been developed to facilitate various graph-related tasks including node classification[16], link prediction[19, 35], and graph classification[36]. Specifically, for graph-related learning tasks, graph pooling has become an essential component in various GNN architectures. Aiming to learn compact representation for graphs, graph pooling facilitates graph topology modeling by enlarging receptive fields as well as scaling down the graph size which effectively reduces computational costs.

The existing graph pooling techniques generally fall into two main categories, i.e., the global graph 29 pooling [5, 27, 14, 33, 26, 36] and hierarchical graph pooling. The former directly compresses a set 30 of nodes into a compact graph-level representation. This operation results in a flat feature as a whole 31 graph embedding. In contrast, the hierarchical pooling coarsens graphs gradually and outputs the 32 corresponding pooled graphs of smaller sizes. For this purpose, two different types of coarsening, 33 named node dropping [17, 10, 20, 37, 18, 11] and node clustering [33, 1, 34], are often employed. 34 The node dropping picks up a subset of nodes to construct the coarsened graph, while the node 35 clustering learns an assignment matrix to aggregate those nodes in the original graph into new clusters. 36 In this work, our proposed BernPool falls in the category of the latter. 37

Even though considerable progress has been made, most pooling methods select a part of nodes or 38

cluster nodes in a deterministic manner according to the importance scores of nodes, which degrades 39

40

the sampling diversity. For this issue, some previous works [21, 7] propose stochastic node dropping

independent from the data, i.e., randomly dropping. However, they do not consider the intrinsic 41

structural characteristics of data while enriching the sampling diversity. Thus the current bottleneck 42 is how to adaptively extract expressive substructures and meanwhile keeping rich sampling diversity 43 in graph pooling, with the precondition of high-efficient and effective graph representation learning. 44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

To address the above problem, in this work, we propose a graph Bernoulli pooling method called BernPool to facilitate graph representation learning. Different from the existing graph pooling methods, we design a probabilistic Bernoulli sampling by estimating the sampling probabilities of graph nodes. To restrict the sampling process, we formulate a variational Bernoulli learning constraint by deriving an upper bound between an expected distribution and a learned distribution. To better capture expressive info, a learnable reference set is further introduced to encode nodes into a latent expressive probability space. Thus the advantage of the resultant Bernoulli sampling is two-fold: i) capture representative substructures of graph; and ii) preserve certain diversity (like random dropping) due to its non-deterministic manner. Considering the complementary characteristics between node dropping and node clustering, we propose a hybrid graph pooling paradigm to fuse a compact subgraph after node dropping and a coarsening graph after node clustering. The node clustering in our framework is also Bernoulli-induced without high-computation cost because it adopts the sampled nodes as clustering centers. The hybrid graph pooling can jointly learn representative substructures and preserve the input graph topology. We conduct extensive experiments on 8 public graph classification datasets to test our BernPool, and the experimental results validate that our BernPool achieves better performance than those existing pooling methods and keep high efficiency on par with those node-dropping methods.

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

The contributions of this work are summarized as: i) propose a probabilistic Bernoulli sampling method to not only learn effective sampling but also preserve high efficiency; ii) propose a hybrid graph pooling way to retain both those sampled substructures and the remaining info; iii) verify the effectiveness and high-efficiency of our BernPool, and report the state-of-the-art performance.

2 Related Work

In this section, we first review the previous methods of Graph Neural Networks (GNNs), then introduce the related Hierarchical pooling methods.

Graph Neural Network. GNNs were introduced as a form of recurrent neural network by Gori et.al. [12] and Scarselli et al. [23]. Subsequently, Duvenaud et al.[6] introduced a convolutionlike propagation rule on graphs to extract node representations for graph-level classification. To enhance the graph representation ability, several convolution operations were proposed (e.g. Graph Convolution Network(GCN[16]), Graph Attention Network(GAT[25]), GraphSAGE[13], GIN[30]) to extract expressive node representations by aggregating neighbor node features and have achieved promising performance in various graph-related tasks in recent years. In particular, GCN[16] utilizes a first-order approximation of spectral convolution via Chebyshev polynomial iteration to improve efficiency. However, it suffers from the issue of fixed and equal weighting for neighbor nodes during aggregation, which may not be optimal for all nodes and can lead to information loss. To address this problem, GAT[25] introduced an attention mechanism to assign different weights to neighbor nodes during message passing. Furthermore, GrapSAGE[13] learned node embeddings by aggregating feature information from the neighborhood in an inductive manner. Despite the considerable progress made by GNNs, they are limited in their ability to generate hierarchical graph representations due to the lack of pooling operations.

Hierarchical Graph Pooling. Graph pooling is a critical operation to obtain robust representations and scale down the size of graphs, which can be classified into two categories: global pooling and hierarchical pooling. The former [5, 27, 14, 33, 26, 36] aggregates node-level features to generate a graph-level representation. For instance, SortPool[36] ranks and groups the nodes into clusters according to their features, then aggregates the resulting clusters to generate the graphlevel representation. However, the global pooling suffers from the issue of discarding structure information in generating graph-level representation. On the other hand, the hierarchical pooling methods could progressively compress the graph into a smaller one and capture the hierarchical structure. They can be further divided into node clustering pooling methods [33, 1, 34], node drop pooling methods [17, 10, 20, 37, 18, 11] and other pooling methods[28, 3]. Among them, the node drop pooling methods deleted the unimportant nodes based on certain criteria. For instance, SAGPool [17] computed the node attention scores using graph convolution to preserve the most

important nodes. But the node drop pooling methods may not preserve the original structure well 96 during the graph compression process. To alleviate this problem, Pang et.al [20] applied contrastive 97 learning to maximize the mutual information between the input graph and the pooled graphs to 98 preserve the graph-level dependencies in the pooling layers. Gao et.al [11] proposed a criterion to 99 assess each node's information among its neighbors to retain informative features. Our BernPool 100 introduces probabilistic deduced Bernoulli sampling based on reference set to progressively compress 101 the original graph by preserving important nodes, rather than selecting nodes in a deterministic 102 way. This probabilistic manner can lead to more diverse sampling situations and capture the data of 103 intrinsic characteristics, promoting graph discriminative representation learning. Furthermore, we 104 propose a hybrid graph pooling module to alleviate the node drop pooling method's issue of cannot 105 preserve structure well. 106

107

108

109

110

111

112

113

114

115

116

117

118

119

3 Preliminaries

Notations For an arbitrary graph G = ( V , E , X ) with n = |V| nodes and |E| edges. X ∈ R n × d ′ represents the node feature matrix, where d ′ is the dimension of node attributes, and A ∈ R n × n denotes the adjacency matrix describing its edge connection information. The graph G has a one-hot vector y i w.r.t its label. A pooled graph of the original graph G is denoted by ˜ G = ( ˜ V , ˜ E , ˜ X ) with adjacency matrix as ˜ A ∈ R ˜ n × ˜ n , where ˜ n denotes the number of pooled nodes.

Graph Convolution In this work, we employ the classic graph convolution network (GCN) as the backbone to extract features, where the l -th convolutional layer is formulated as:

where σ ( · ) is a non-linear activation function, H ( l ) is the hidden-layer feature, ˆ A is the added self-loop adjacent matrix, ˆ D denotes the degree matrix of ˆ A , and W ( l ) represents a learnable weight matrix at the l -th layer. The initial node features are used at the first convolution, i.e., H (0) = X .

4 The Proposed BernPool

4.1 Overview

Graph

Conv

Graph

Conv

Graph

Conv

BernPool

BernPool

BernPool

Readout

Readout

Readout

MLP

Reference Set

Reparametrization

Predicted

Label

Importance score

Pooled graph

Hybrid Graph Pooling

dropping

clustering

g = (X,A)

O

(H,A)

0.82

p

0.68

0

B(p)

.".

0.78

1

(H,A)

0.85

1

S

(HI,A)

The whole framework is illustrated in Fig. 1, where the convolution and our proposed BernPool are 120 stacked alternately. The proposed BernPool could be seamlessly engaged with any type of graph 121 convolution to facilitate graph representation. In BernPool, there contains two main modules: graph 122 Bernoulli sampling (GBS) and hybrid graph pooling (HGP). To boost graph Bernoulli sampling, 123 we specifically design a reference set S to encode the importance of each graph node, and the 124 reference set is configured as the optimizable parameters. The reference set is further conditioned on 125 the orthogonal space so as to reduce redundancy. After transforming via the reference set, we can 126

estimate the sampling probabilities p of graph nodes. In particular, the probabilities of nodes are 127 totally restricted with an expected/predefined distribution, which is formulated to maximize the upper 128 bound of KL-divergence (please see Section 4.2). Just due to the restriction, we can probabilistically 129 sample a specified proportion of graph nodes. The detail of GBS can be found in Section 4.3. In the 130 stage of hybrid graph pooling, on the one hand, we prune those unsampled nodes and the associated 131 edges to generate a compact subgraph; on the other hand, to preserve graph topological structure, 132 we perform neighbor nodes clustering to form a coarsening graph. Both the compact subgraph 133 and coarsening graph are fused to form the final pooled graph. The detail of HGP can be found in 134 Section 4.4. The BernPool attempts to learn the reference set and a few linear transformations. The 135 whole framework can be optimized in an end-to-end mode through back-propagation. 136

137

138

139

140

141

4.2 Bernoulli Sampling Optimization Objective

To sample a certain proportion of graph nodes in a probabilistic manner, we derive a KL-divergence constraint, which makes node probabilities tend to be a predefined distribution. To this end, we again dissect mutual information between the learned subgraph embeddings and their corresponding labels. Formally, we aim to maximize the mutual information function:

where ϕ denotes the parameters of the BernPool, ψ is the parameters of other modules (e.g., convolu142 tion, classifier), and f ψ,ϕ ( · ) represents the graph embedding process. 143

Suppose the sampling factor z in graph pooling, we resort to the relationship between mutual 144 information and expectation, and rewrite Eqn. (2) as: 145

where p ψ ( y |G , S , z ) is the conditional probability of label y , and p ϕ ( z |G , S ) denotes the conditional 146 probability of the factor z that is usually intractable. After a series of derivation from Eqn. (3), we 147 can deduce a bound with KL-divergence between expected Bernoulli distribution q ϕ and learned 148 distribution p ϕ : 149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

where p ψ ( y |G , S ) represents the predicted probability of label based on the input graph and reference set, ζ CE is the cross entropy loss function. Please see the detailed derivation in the supplementary file . As the expected Bernoulli distribution q ϕ ( z |G , S ) is independent from the input graph and reference set, q ϕ ( z |G , S ) can be denoted as q ( z ) . After adding the soft-orthogonal constraint on the reference set, therefore, the final optimization objective can be converted to minimize:

where the second term forces the sampling factor to follow an expected distribution q , the matrix S stacks the vectors of reference set S in the third term, I denotes the identity matrix, β is a trade-off hyper-parameter, and c is a learnable scalar. Thus, the learning of the sampling factor could be integrated into the objective function as a joint training process. In addition, we can easily extend the above single-layer BernPool into multi-layer networks by deploying independent sampling factors in sequential graph pooling.

4.3 Bernoulli Sampling Factor Learning

To extract an expressive sub-graph ˜ G from the original or former-layer graph G , we estimate a probabilistic factor z = ( z 1 , · · · , z n ) ⊺ ∈ { 0 , 1 } n that conforms to Bernoulli distribution, instead of a deterministic way such as top-k. In contrast to the deterministic way, our BernPool possesses

a more diverse sampling in mining substructures and graph topological variation of input data. 165 However, as mentioned above in Eqn. (3), it's rather non-trivial to infer z through Bayes rule: 166 p ( z |G , S ) = p ( z ) p ( G , S| z ) /p ( G , S ) . A reason is that the prior p ( G , S| z ) is intractable. We resort 167 to the variational inference to approximate the intractable true posterior p ( z |G , S ) with q ( z ) by 168 constraining the KL-divergence D KL ( q ϕ ( z ) || p ϕ ( z |G , S )) . 169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

To make the gradient computable, we employ the reparameterization trick to derive p ( z |G , S ) = B ( p ) , where B ( · ) denotes the Bernoulli distribution based on the probability vector. Concretely, we take the graph convolution feature H and the matrix S of reference set as input, to learn the sampling probability p . A simple formulation with a fully-connected operation is given as follows:

where cos( H , S ) computes the pairwise similarities (across all graph nodes and all reference points) through the cosine measurement, and sigmoid ( · ) is the sigmoid function.

4.4 Bernoulli Hybrid Graph Pooling

Based on the above inferred sampling factor z , we propose a hybrid graph pooling to learn expressive substructures while endeavoring to reserve graph information. The hybrid graph pooling contains two components: Bernoulli node dropping and Bernoulli node clustering. The former prunes those unsampled nodes and associated edges to generate a compact graph; the latter clusters neighbor nodes to form a coarsening graph.

Bernoulli Node Dropping. Let idx ∈ R ˜ n denotes the indices of the preserved nodes according to the sampling factor z , thus a projection matrix P ∈ { 0 , 1 } ˜ n × n can be defined formally:

where diag ( · ) is the vector diagonalization operation, and X [idx , :] extracts those rows of X w.r.t the indices idx . Accordingly, the compact subgraph ˜ G = ( ˜ H , ˜ A ) can then be computed by:

where ˜ H ∈ R ˜ n × d is the feature matrix of the compact graph, and ˜ A ∈ R ˜ n × ˜ n is the subgraph adjacency matrix. Intuitively, the dropping directly removes those unselected nodes and connective edges from the input graph.

Bernoulli Node Clustering. After Bernoulli sampling, we can obtain those reserved nodes, which could be used as the clustering centers. Hereby, we only need to transmit those unselected nodes' messages to cluster centers, which would preserve more information of the whole input graph. Compared with previous graph clustering methods, Bernoulli clustering is more efficient because the learning of the assignment matrix is bypassed, while increasing the diversity of graph perception. To be specific, our assignment matrix P ′ ∈ R ˜ n × n is from the original adjacency matrix A , and we can get the coarsening graph as:

where ˜ H ′ ∈ R ˜ n × d denotes the diffusion features based on input graph, the coarsening graph has the same adjacency matrix ˜ A i as in Eqn. (9).

Hybrid Graph. Based on the aforementioned operations, we can obtain the compact subgraph ˜ G = ( ˜ H , ˜ A ) and the coarsening graph ˜ G ′ = ( ˜ H ′ , ˜ A ) , which reflect different aspects of information in the original graph. To fully exploit the extracted informative node representations and expressive substructures, the two subgraphs are fused to form the final pooled graph by the following formulation:

where σ denotes a non-linear activation function and W h ∈ R d × d is a learnable weight, and ̂ H ∈ R ˜ n × d is the aggregated node embeddings of the pooled graph.

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

4.5 Readout Function

The proposed framework repeats the graph convolution and BernPool operations three times. To obtain a fixed-size graph-level representation, we apply the concatenation of max-pooling and meanpooling in each subgraph following the previous works [37, 29, 22]. Finally, those graph-level representations can be summarized to form the final embeddings:

where ̂ H ( l ) i denotes the i -th node feature at the l -th pooling, σ is the same non-linear activation function, and || denotes the feature concatenation operation. The resulting embeddings would finally be fed into a multi-layer perceptron to predict graph labels.

4.6 Computational Complexity

The computational complexity of one-layer BernPool can be expressed as O ( N × K × d + N ′ × N × d + N ′ × d × d ) , where N denotes the number of nodes, K is the number of reference points, d is the dimensionality of nodes features, N ′ represents the number of preserved nodes. Specifically, the complexity of probability score computation is O ( N × K × d ) , and the complexity of the Bernoulli hybrid graph pooling module is O ( N ′ × N × d + N ′ × d × d ) .

5 Experiments

5.1 Experimental Setup

Datasets . To comprehensively evaluate our proposed model, we conduct extensive experiments on eight widely used datasets in the graph classification task, including three social network datasets (IMDB-BINARY, IMDB-MULTI [31] and COLLAB[32]) and five Bioinformatics datasets (PROTEINS[8], DD[4], NCI1[24], Mutagenicity[15] and ENZYMES[2]). The detailed information and statistics of these datasets are summarized in Table. 1.

Baselines . We compare our proposed method with several state-of-the-art graph pooling methods, including three backbones (GCN, GAT, GraphSAGE), eight node drop graph pooling methods (TopkPool[10], SAGPool[17], ASAP[22], VIPool[18], iPool[11], CGIPool[20], SEP-G [29] and MVPool [37]), three clustering pooling methods (Diffpool[33], MincutPool[1], and StructPool[34]), three global pooling methods (Set2Set[26], SortPool[36], DropGIN[21]) and one other pooling method (EdgePool[3]).

Implementation Details . We employ the 10-fold cross-validation protocol following the settings of [29, 22] and report the average classification accuracies and standard deviation. For all used datasets, we set the expected pooling ratio as 0.8, the node embedding dimension d as 128, the number of reference points as 32, and the hyper-parameter β in Eqn. 6 as 5. We adopt the Adam optimizer to train our model with 1000 epochs, where the learning rate and weight decay are set as 1e-3 and 1e-4, respectively. Our proposed BernPool is implemented with PyTorch and Pytorch Geometric [9].

5.2 Comparison with the state-of-the-art Methods

The graph classification results of BernPool and other state-of-the-art methods are presented in Table 1. In general, most hierarchical pooling approaches including our proposed BernPool can perform better than those global pooling ones in the graph classification task. This may be because global pooling methods ignore the hierarchical graph structures in generating graph-level representation. In particular, our BernPool achieves state-of-the-art performance on all datasets, which demonstrates the robustness of our framework against graph structure data variation. In contrast, previous methods cannot perform well on all eight datasets, while the second highest performances on different datasets are obtained by five different methods. Compared with those methods, Our BernPool outperforms respectively by 1.09%, 3.16%, 13.6%, 2.7%, and 1.86% on the PROTEINS, Mutagenicity, ENZYMES, IMDBBINARY and COLLAB datasets.

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

Moreover, the proposed BernPool significantly surpasses GNNs (GCN, GAT, GraphSAGE) without adding the pooling operation, which verifies our BernPool can effectively learn representative substructures and preserve graph topological information. SEP-G [29] gains better performance compared with other node drop pooling methods, which utilizes the structural entropy to assess the importance of each graph node. However, our BernPool exhibits an average 5% relative improvement over SEP-G. This can be attributed to our method of learning sampling factors in a probabilistic manner, which possesses more diversity in mining expressive sub-structures compared with the deterministic way. Compared to DiffPool [33], our BernPool outperforms it on all used datasets, verifying the effectiveness of our proposed hybrid graph pooling module, which jointly leverages the advantages of both node drop and clustering methods. Particularly, DropGIN [21] randomly drops nodes in the training process, which is also a probabilistic manner. Our BernPool outperforms it by 6.99%, 5.6%, and 4.53% on the PROTEINS, IMDB-BINARY, and IMDB-MULTI datasets. This is because our BernPool considers the characteristics of data and can adaptively handle the topology variations.

6 Ablation Study

Effectiveness of the proposed BernPool using GCN variants. To evaluate the performance of our method by employing different convolution layers, we integrate three widely used ones (i.e., GCN[16], GAT[25] and GraphSAGE[13]) into our BernPool framework. The results on eight datasets are shown in Table 2. It can be observed that all three variants (BernPool-GCN, BernPool-GAT, BernPool-GraphSAGE) outperform their corresponding backbones, significantly improving the graph classification performances. This observation validates the effectiveness of the proposed BernPool.

Effectiveness of Bernoulli hybrid graph pooling. To verify the effectiveness of our proposed Bernoulli hybrid graph pooling, we further conduct experiments on eight datasets. As the hybrid graph pooling consists of node dropping and clustering. We separately remove one of the channels and keep the other parts the same. We name the BernPool without dropping and clustering as "BernPool w/o Dropping" and "BernPool w/o Clustering", respectively. The detailed results are reported in Table 3.

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

Notably, our BernPool employs just one channel can achieve good results, which demonstrates the effectiveness of BernPool leveraging a probabilistic manner to infer sampling factors. We can observe that the "BernPool w/o Dropping" outperforms "BernPool w/o Clustering" by 0.46%, 0.42%, 0.70%, and 0.54% on PROTEINS, DD, IMDB-BINARY, and COLLAB datasets, respectively. Furthermore, jointly using both channels can outperform either "BernPool w/o Dropping" or "BernPool w/o Clustering", which verifies the effectiveness of our proposed hybrid graph pooling.

Comparison between the probabilistic and deterministic manner. To make clear the benefit of our proposed probabilistic sampling method, we conduct experiments by replacing Bernoulli sampling with Topk pooling which is in a deterministic manner. In the TopK experiments, we still employ reference set to assess the importance of nodes and the hybrid graph pooling module to jointly learn representative sub-structures and preserve graph topological information. The comparison between "BernPool-TopK" and "BernPool" is presented in Table 4. It can be observed that BernPool outperforms BernPool-TopK on all used eight datasets, especially 7.83% accuracy gains in the ENZYMES dataset, verifying the effectiveness of our designed Bernoulli-deduced sampling strategy.

Benefit of the orthogonality for reference set. To evaluate the benefit of orthogonality for the reference set, we conduct experiments that remove the orthogonal constraint from BernPool (referred to as "BernPool w/o orthogonal" in Fig. 2(a)). The results demonstrate that employing the orthogonal reference set can improve average accuracy by more than 0.6% on three datasets. This verifies the effectiveness of the orthogonal constraint for the reference set.

Model parameter quantity comparison. We compare the test accuracy and parameter quantity (only the pooling layer) of BernPool with other pooling methods in the PROTEINS dataset, where the hidden layer dimension is set to 128. BernPool achieves superior performance while using fewer parameters. Specifically, BernPool owns 97% fewer parameters than CGIPool and 76% fewer parameters than ASAP. In terms of accuracy, BernPool performs the best, achieving 9.19% higher than CGIPool and 9.10% higher accuracy than ASAP.

IMDB-B

PROTEINS

COLLAB

60

65

70

75

80

85

81.0

82.22

84.7

81.3

83.29

85.26

BernPool w/o orthogonality

BernPool

(a) Benefit of orthogonality for ref-

0

100

200

300

400

500

70

72

74

76

78

80

82

BernPool

TopkPool

ASAP

EdgePool

SAGPool

CGIPool

(b) Parameter quantity (KB)

0.0

0.2

0.4

0.6

0.8

70

72

74

76

78

80

82

BernPool

TopkPool

ASAP

EdgePool

SAGPool

CGIPool

(c) Inference time (seconds)

erence set

0.2

0.3

0.4

0.5

0.6

0.7

0.8

50

55

60

65

70

75

80

85

PROTEINS

NCI1

(a) Pooling ratio

1

2

3

4

50

55

60

65

70

75

80

85

PROTEINS

NCI1

(b) Model layers

32

64

128

256

50

55

60

65

70

75

80

85

PROTEINS

NCI1

(c) Hidden layer dimension

Inference time comparison. The computation complexity of our BernPool is described in Section 4.6. 299 Moreover, the inference time also plays a crucial role in evaluating the efficiency of pooling methods. 300 Thus we conduct experiments on the PROTEINS dataset containing about 20,000 nodes and 70,000 301 edges to compare the inference time (single-layer). As shown in Fig. 2(c), BernPool costs less 302 inference time than EdgePool and ASAP while maintaining superior performance. Notably, SAGPool 303 has a similar inference time as our method, but our method outperforms it in terms of classification 304 accuracy. The comparison results verify the high efficiency of our BernPool. 305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

Sensitivity of Hyper-parameters. To evaluate the sensitivity of hyper-parameters, including the pooling ratio, layer number, and hidden layer dimension, we additionally conduct experiments on PROTEINS and NCI1 datasets. Specifically, we vary the pooling ratio from 0.2 to 0.8 with a step length of 0.2, the number of layers from one to four, and the hidden layer dimensions range from 16 to 128. The results are presented in Fig. 3. We can observe that BernPool overall exhibits robustness to variations of parameters. However, performance fluctuations can be observed on the NCI1 dataset when we evaluate the effect of different pooling ratios for BernPool. This may be attributed to the relatively less average number of nodes resulting in less information being retained after performing three pooling layers consecutively. The results shown in Fig. 3(b) indicate that setting the layer number to three achieves the best performance on the PROTEINS and NCI1 datasets. However, increasing the number of layers will require more computation resources and longer training time. As shown in Fig. 3(c), the highest accuracy is achieved when the dimension size is set as 128. With the dimension increasing, the accuracy presents a slight increase trend, which suggests that increasing the dimension size can enhance the model's capacity to capture more complex representations of the input graph. However, larger dimension sizes increase computation burdens. Thus we set the pooling ratio as 0.8, the layer number as 3, and the dimension as 128 in our framework.

322

323

324

325

326

327

328

329

330

331

332

7 Conclusion

In this paper, we proposed a simple and effective graph pooling method, called Graph Bernoulli Pooling (BernPool) to promote the graph classification task. Specifically, a probabilistic Bernoulli sampling was designed to estimate the sampling probabilities of graph nodes, and to further extract more useful information, we introduced a learnable reference set to encode nodes into a latent expressive probability space. Compared with the deterministic way, BernPool possessed more diversity to capture salient substructures. Then, to jointly learn representative substructures and preserve graph topology information, we proposed a hybrid graph pooling paradigm that fuses two pooling manners. We evaluate BernPool on multiple widely used datasets and dissected the framework with ablation analysis. The experimental results show that BernPool outperforms state-of-the-art methods and demonstrates the effectiveness of our proposed modules.

333

334

335

336

337

338

339

References

Filippo Maria Bianchi, Daniele Grattarola, and Cesare Alippi. Spectral clustering with graph neural networks for graph pooling. In International conference on machine learning , pages 874-883. PMLR, 2020.

Karsten M Borgwardt, Cheng Soon Ong, Stefan Schönauer, SVN Vishwanathan, Alex J Smola, and Hans-Peter Kriegel. Protein function prediction via graph kernels. Bioinformatics , 21(suppl_1):i47-i56, 2005.

Frederik Diehl. Edge contraction pooling for graph neural networks. arXiv preprint 340 arXiv:1905.10990 , 2019. 341

Paul D Dobson and Andrew J Doig. Distinguishing enzyme structures from non-enzymes 342 without alignments. Journal of molecular biology , 330(4):771-783, 2003. 343

David K Duvenaud, Dougal Maclaurin, Jorge Iparraguirre, Rafael Bombarell, Timothy Hirzel, 344 Alán Aspuru-Guzik, and Ryan P Adams. Convolutional networks on graphs for learning 345 molecular fingerprints. Advances in neural information processing systems , 28, 2015. 346

David K Duvenaud, Dougal Maclaurin, Jorge Iparraguirre, Rafael Bombarell, Timothy Hirzel, 347 Alan Aspuru-Guzik, and Ryan P Adams. Convolutional networks on graphs for learning 348 molecular fingerprints. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, 349 editors, Advances in Neural Information Processing Systems , volume 28. Curran Associates, 350 Inc., 2015. 351

352

353

354

Wenzheng Feng, Jie Zhang, Yuxiao Dong, Yu Han, Huanbo Luan, Qian Xu, Qiang Yang, Evgeny Kharlamov, and Jie Tang. Graph random neural networks for semi-supervised learning on graphs. Advances in neural information processing systems , 33:22092-22103, 2020.

355

356

357

358

359

360

361

362

363

364

Aasa Feragen, Niklas Kasenburg, Jens Petersen, Marleen de Bruijne, and Karsten Borgwardt. Scalable kernels for graphs with continuous attributes. Advances in neural information processing systems , 26, 2013.

Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorch geometric. arXiv preprint arXiv:1903.02428 , 2019.

Hongyang Gao and Shuiwang Ji. Graph u-nets. In international conference on machine learning , pages 2083-2092. PMLR, 2019.

Xing Gao, Wenrui Dai, Chenglin Li, Hongkai Xiong, and Pascal Frossard. ipool-informationbased pooling in hierarchical graph neural networks. IEEE Transactions on Neural Networks and Learning Systems , 33(9):5032-5044, 2021.

M. Gori, G. Monfardini, and F. Scarselli. A new model for learning in graph domains. In 365 Proceedings. 2005 IEEE International Joint Conference on Neural Networks, 2005. , volume 2, 366 pages 729-734 vol. 2, 2005. 367

Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large 368 graphs. Advances in neural information processing systems , 30, 2017. 369

Takeshi D Itoh, Takatomi Kubo, and Kazushi Ikeda. Multi-level attention pooling for graph 370 neural networks: Unifying graph representations with multiple localities. Neural Networks , 371 145:356-373, 2022. 372

Jeroen Kazius, Ross McGuire, and Roberta Bursi. Derivation and validation of toxicophores for 373 mutagenicity prediction. Journal of medicinal chemistry , 48(1):312-320, 2005. 374

Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional 375 networks. arXiv preprint arXiv:1609.02907 , 2016. 376

Junhyun Lee, Inyeop Lee, and Jaewoo Kang. Self-attention graph pooling. In International 377 conference on machine learning , pages 3734-3743. PMLR, 2019. 378

Maosen Li, Siheng Chen, Ya Zhang, and Ivor Tsang. Graph cross networks with vertex infomax 379 pooling. Advances in Neural Information Processing Systems , 33:14093-14105, 2020. 380

381

382

383

David Liben-Nowell and Jon Kleinberg. The link prediction problem for social networks. In Proceedings of the twelfth international conference on Information and knowledge management , pages 556-559, 2003.

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

Yunsheng Pang, Yunxiang Zhao, and Dongsheng Li. Graph pooling via coarsened graph infomax. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval , pages 2177-2181, 2021.

Pál András Papp, Karolis Martinkus, Lukas Faber, and Roger Wattenhofer. Dropgnn: Random dropouts increase the expressiveness of graph neural networks. Advances in Neural Information Processing Systems , 34:21997-22009, 2021.

Ekagra Ranjan, Soumya Sanyal, and Partha Talukdar. Asap: Adaptive structure aware pooling for learning hierarchical graph representations. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 5470-5477, 2020.

Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions on Neural Networks , 20(1):61-80, 2009.

Nino Shervashidze, Pascal Schweitzer, Erik Jan Van Leeuwen, Kurt Mehlhorn, and Karsten M Borgwardt. Weisfeiler-lehman graph kernels. Journal of Machine Learning Research , 12(9), 2011.

Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, Yoshua Bengio, et al. Graph attention networks. stat , 1050(20):10-48550, 2017.

Oriol Vinyals, Samy Bengio, and Manjunath Kudlur. Order matters: Sequence to sequence for sets. arXiv preprint arXiv:1511.06391 , 2015.

Zhengyang Wang and Shuiwang Ji. Second-order pooling for graph neural networks. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2020.

Lanning Wei, Huan Zhao, Quanming Yao, and Zhiqiang He. Pooling architecture search for graph classification. In Proceedings of the 30th ACM International Conference on Information &Knowledge Management , pages 2091-2100, 2021.

Junran Wu, Xueyuan Chen, Ke Xu, and Shangzhe Li. Structural entropy guided graph hierarchical pooling. In International Conference on Machine Learning , pages 24017-24030. PMLR, 2022.

Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint arXiv:1810.00826 , 2018.

Pinar Yanardag and SVN Vishwanathan. Deep graph kernels. In Proceedings of the 21th ACM SIGKDD international conference on knowledge discovery and data mining , pages 1365-1374, 2015.

Pinar Yanardag and SVN Vishwanathan. A structural smoothing framework for robust graph comparison. Advances in neural information processing systems , 28, 2015.

Zhitao Ying, Jiaxuan You, Christopher Morris, Xiang Ren, Will Hamilton, and Jure Leskovec. Hierarchical graph representation learning with differentiable pooling. Advances in neural information processing systems , 31, 2018.

Hao Yuan and Shuiwang Ji. Structpool: Structured graph pooling via conditional random fields. In Proceedings of the 8th International Conference on Learning Representations , 2020.

Muhan Zhang and Yixin Chen. Link prediction based on graph neural networks. Advances in neural information processing systems , 31, 2018.

Muhan Zhang, Zhicheng Cui, Marion Neumann, and Yixin Chen. An end-to-end deep learning 424 architecture for graph classification. In Proceedings of the AAAI conference on artificial 425 intelligence , volume 32, 2018. 426

427

Zhen Zhang, Jiajun Bu, Martin Ester, Jianfeng Zhang, Zhao Li, Chengwei Yao, Huifen Dai,

428

429

Zhi Yu, and Can Wang. Hierarchical multi-view graph pooling with structure learning. IEEE Transactions on Knowledge and Data Engineering , 35(1):545-559, 2021.

Peer reviews:
comments: SUMMARY:
This paper studies the problem of graph classification and proposes a method named graph Bernoulli pooling (BernPool).  Considering the complementarity of node dropping and node clustering, the paper proposes a hybrid graph pooling paradigm to combine a compact subgraph (via dropping) and a coarsening graph (via clustering), in order to retain both representative substructures and input graph info.

STRENGTHS:
1. The problem is important.

2. The writing is good.

WEAKNESSES:
1. The results are weirdly high. I checked the code and found no validation set. First, the experimental setting is different from MVGPool, but author directly copies the number. Second, for example, authors select the worst model variant of MVGPool in the original paper. 

2. Node dropping and clustering are both well-known strategies for graph pooling. Combining them is not novel to this community. 

3. The ablation studies show your both variants highly outperform the best baseline on ENZYME, for example. So, which part works most in your framework, not the clustering or dropping? Can you explain more?

QUESTIONS:
see weaknesses.

LIMITATIONS:
see weaknesses.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 1
DATE: 2023-06-08

comments: SUMMARY:
Graph pooling is an essential operator in GNNs for graph classification. This paper propose a hierarchical pooling method which is based on the node clustering. Existing graph pooling methods are designed in a deterministic way which neglect the intrinsic structure feature. This paper propose BernPool, which first conducts Bernoulli sampling and then combines a compact subgraph after node dropping and a coarsening graph after node clustering. BernPool can jointly learn expressive substructures and keep graph topology. The core of this paper is the utilization of Bernoulli sampling factor for node selection and clustering.

STRENGTHS:
(1)	The topic of graph pooling is important in graph representation learning.

(2)	The writing of this paper is clear.

WEAKNESSES:
(1)	The model's complexity does not effectively demonstrate the necessity or unique advantages of each intricate design choice. The model first adopts a complex strategy for sampling, and then employs a complex hybrid pooling strategy, which lacks significant insights into their effectiveness.

(2)	Insufficient experimentation. The dataset used in the experiments is too small, and the effectiveness should be validated on larger datasets, such as ogb and zinc.

(3)	The method has excessively high complexity, which can limit its scalability in practical applications.

QUESTIONS:
Does the proposed method have higher expressive power?

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a probabilistic hierarchical graph pooling method called graph Bernoulli pooling (BernPool) for graph features learning.\nUnlike previous deterministic works, the probabilistic Bernoulli sampling with the reparameterization trick allows BernPool to preserve certain diversity during the downsampling without losing high efficiency.\nBesides, BernPool can capture representative substructures of graphs by combining the advantage of node-dropping and node-clustering techniques.\nAmple empirical experiments demonstrate the effectiveness of the proposed methods, which outperform previous works with remarkable gaps and fewer parameters.\nThe ablation study justifies each proposed component in BernPool.\n\nSTRENGTHS:\n1. This work proposes new a hierarchical pooling method combining node-clustering and node-dropping and introduces probabilistic properties with a Bernoulli sampling framework. The Bernoulli optimization objective with soft-orthogonal constraint also contributes to performance improvement.\n2. The proposed method outperforms previous techniques with remarkable performance gaps. Besides, the method is also applicable to several typical GNN backbones and shows decent improvements.\n3. The proposed method is well supported by comprehensive empirical comparisons and ablation studies.\n4. The mathematical derivation supports the proposed Bernoulli sampling well theoretically.\n\nWEAKNESSES:\n1. There are some minor confusing points in the paper writing. (listed in the question)\n2. The authors explain the motivation for using probabilistic techniques in selecting reference sets and mention the drawbacks of node-dropping pooling methods. \nHowever, the motivation for the hybrid graph pooling is vague. The authors categorize BernPool as a node clustering pooling method but do not explain the motivation for adding back node dropping counterpart.\n\nQUESTIONS:\n1. I am a bit confused about the reparameterization trick mentioned in line 170 for Bernoulli distribution. I assume it shall be something similar to [1]. If so, I think you should cite it. If I am wrong or missing something, please provide more details.\n2. What is the motivation for adding node-dropping, and why it helps? \nAccording to the ablation study, removing node-dropping will lead to a slight performance drop except on NCI1 and ENZYMES, on which BernPool is impacted notably.\n3. How BernPool works during the inference? Especially regarding the Bernoulli sampling. Will BernPool have multiple inferences with different sampling, similar to MC-dropout [2] and the reference [21] listed in the paper? \n4. When introducing multiple BernPool layers in a neural network,  $\\mathbf{H}$ in eq. (7) will be the same for all BernPool layers, or it refers to the node representations of the coarsened graphs at the corresponding layer?\n\n\n\n------------\n- [1] Bengio, Yoshua, Nicholas Léonard, and Aaron Courville. \"Estimating or propagating gradients through stochastic neurons for conditional computation.\" arXiv preprint arXiv:1308.3432 (2013).\n- [2] Gal, Yarin, and Zoubin Ghahramani. \"Dropout as a Bayesian approximation: Representing model uncertainty in deep learning.\" international conference on machine learning. PMLR, 2016.\n\nLIMITATIONS:\nNo", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
This paper proposes a pooling method that probabilistically generates a coarse graph using an expected sampling strategy. A hybrid pooling paradigm is designed to integrate two existing pooling paradigms. Extensive experiments demonstrate the effectiveness of the proposed method.

STRENGTHS:
The method is well-organized and easy to follow.

The proposed approach is simple yet effective, as supported by the experimental results.

WEAKNESSES:
The stability of the proposed method is unclear. Since different runs may result in different factors z, which in turn lead to different projection matrices P and coarse graphs, the stability of the method needs to be considered. It would be important to address this issue, as obtaining different results in each inference procedure can be problematic.

In the experiments, it would be beneficial to provide examples that illustrate the reference set S, probability P, and the factor z. Including these examples would enhance the readers' understanding of the proposed method and its components.

QUESTIONS:
Please check the weaknesses.

LIMITATIONS:
Please check the weaknesses.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors introduce a non-deterministic graph pooling method called graph Bernoulli pooling (BernPool). This work comprises two main components: a probabilistic Bernoulli sampling method that allows for more diverse sampling situations and captures intrinsic characteristics of the data, and a hybrid graph pooling paradigm that combines a compact subgraph with a coarsened graph. Experimental results demonstrate the significant superiority of BernPool over the baseline methods in graph classification tasks.\n\nSTRENGTHS:\nThis work is easy to follow.\nThe idea presented in this paper is straightforward and the hybrid graph pooling technique significantly improves the performance of graph pooling. \nThe authors have provided code implementation, enhancing reproducibility and adding to the persuasiveness of their work.\n\nWEAKNESSES:\n1. More visualizations are needed to illustrate how node dropping and node clustering work, such as which nodes are discarded during the pooling process and which nodes are grouped into the same clusters.\n2. Transforming the deterministic process of pooling into sampling has been explored in previous works [1]. Although BernPool presents a more refined approach, the authors provide insufficient discussion or theoretical analysis of why Bernoulli sampling is superior to deterministic algorithms.\n[1] Liu, Ning, et al. \"Unsupervised Hierarchical Graph Pooling via Substructure-Sensitive Mutual Information Maximization.\"\n\nQUESTIONS:\nQ1. Regarding the inference time, I guess that in line 301, when the authors mention the number of nodes and edges in the \"proteins\" dataset, it refers to the total number of nodes and edges across all graphs in the dataset (20,000/70,000). However, it is worth noting that the average number of nodes and edges of the graphs in the \"proteins\" dataset is relatively low (approximately 30/70). This might give readers a misconception that BernPool can effectively handle large graphs with tens of thousands of nodes. However, it is important to consider that the scalability of node clustering pooling often not be as efficient as node dropping pooling. The use of dense adjacency matrices in node clustering pooling can result in a higher space complexity of O(n^2), which limits its ability to handle large graphs. To provide further clarity on the contribution of BernPool, the authors could consider running experiments on datasets with a larger number of nodes (e.g., D&D dataset) and provide information on the memory usage and running time of BernPool compared to other baseline methods (e.g., DiffPool, MinCutPool, SAGPool, Topkpool). This would help in better understanding the capabilities and limitations of BernPool.\n\nLIMITATIONS:\nThe proposed method appears to be somewhat incremental.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}