Graph Classification via Reference Distribution Learning: Theory and Practice

Response to Weakness 1:

Thanks for your comment. This is a misunderstanding. Almost all papers studying the theory of classification use the i.i.d assumption. The identical distribution does NOT contradict the multiple classes scenario. Here the identical data distribution $\mathcal{D}$ is actually a composition of $K$ sub-distributions, where each sub-distribution (related to one reference distribution in our method) corresponds to a class. An intuitive example is the Gaussian mixture models: data distribution $\mathcal{D}$ is composed of $K$ Gaussians, i.e., $p_{\mathcal{D}}\left(\mathbf{x}\right)=\sum_{k=1}^K p\left(\mathbf{x} \mid \mathbf{z}\right) p(\mathbf{z})=\sum_{k=1}^K \pi_k \mathcal{N}\left(\mathbf{x} \mid \mu_k, \Sigma_k\right)$ (samples drawn i.i.d from $p_{\mathcal{D}}$ belongs to different classes). For instance, the following popular papers on classification used i.i.d assumption as us.

[1] VN Vapnik. The nature of statistical learning theory. Springer science & business media, 2013. (68000+ citations)

[2] PL Bartlett et al. Spectrally-normalized margin bounds for neural networks. NeurIPS 2017. (1200+ citations)

[3] C. Zhang et al. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM 64, no. 3 (2021): 107-115. (7000+ citations)

Response to Weakness 2:

We treat the nodes' embeddings as discrete distributions and measure the distance between two distributions using MMD. MMD [1] with a Gaussian kernel actually compares all orders of statistics between two distributions rather than the means, though this is implicitly conducted in the RKHS using the mean difference. In other words, MMD compares also the higher-order statistics between distributions while mean pooling compares only the first-order statistics, i.e., mean. Therefore, the information loss of MMD is much less than the mean and max poolings.

We used Example 2.1 to show the limitation of mean and max pooling and the motivation for we regard each node embedding matrix as a discrete distribution. Two different distributions may have the same mean and max but their MMD is never zero. MMD has also been utilized in generative models (e.g. MMD-GAN [2]) to compare distributions, where comparing the means (like mean pooling) in the feature space given by a neural network does not work.

[1] Arthur Gretton et al. A kernel two-sample test. JMLR 2012.

[2] Li et al. MMD GAN: Towards deeper understanding of moment matching network. NeurIPS 2017.

Response to Weakness 3:

We acknowledge that our method is efficient in terms of the number of graphs rather than the number of nodes in each graph. However, there are many graph datasets with a large number of graphs but a relatively small number of nodes in each graph, for which our method can be used.

It is really very difficult to propose an algorithm with both SOTA accuracy and SOTA efficiency. Our method has a good trade-off between them. We follow your suggestion and compare the time cost of our method with two pooling methods published in 2023 (WitTopoPool and MSGNN). The results are in the global rebuttal due to the character limitation. In fact, these two methods are much more costly than ours. The GNN-based topological pooling layer requires the calculation of pairwise node similarity, which is $O(N^2)$. Besides, the time complexity of its witness complex-based topological layer is also quadratic. The subgraph sampling, selection, and evolution of MSGNN are also very costly.

Response to Weakness 4:

Thanks for your insightful comment. Similar to the previous work on learning theory such as [1][2][3], our theorem can only show the impact of the model architecture and parameters on the generalization ability of the model and there is always a trade-off between model complexity and training accuracy, that's why we used the words like "moderate" and "smaller". The upper bounds of the training error are data-dependent, which means we may never point out exactly what model is the best according to only the theoretical analysis.

The generalization bound is the upper bound of the difference between training error and testing error. We say "a network with a smaller $L$ and $r$ may guarantee a tighter bound on..." because in the bound, the term $\tilde{\mathcal{O}}\left(\frac{\mu \bar{b}\|\mathbf{X}\|_2 c^L(L r)^{\frac{3}{2}} \bar{\kappa}^{L r} \sqrt{\theta K / n}}{N}+\gamma \sqrt{\frac{K m d}{N}}\right)$ increases with $L$ and $r$. So here the "smaller" is an exact expression.

We say "moderate-size references" because the bound scales the reference size $m$ as $\tilde{O}(\sqrt{m})$, which means the bound is not very sensitive to $m$ because of the square root. If $m$ is too small, the expressive power of the model will be low; if it is very large, the complexity $\tilde{O}(\sqrt{m})$ is high. Therefore, we say "moderate-size", meaning a trade-off, should be used. This is supported by the experiments in Appendix D.3 (Figure 6) in our paper.

Let's use an intuitive example (may not be true in practice) to further explain why we have to use "moderate" rather than a concrete value (e.g. 0, 1, $\infty$) in a discussion of the theoretical result. Suppose there is a parameter or hyperparameter $s$ of the model and the bounds scales with it as $\tilde{O}((s-1)^2)$ and $s$ does not influence the training error, then the best $s$ is 1. However, in practice, the training error may decrease when $s$ increases. The training error is data-dependent, which means we cannot find the best $s$ using the theoretical result only. Instead, one may use cross-validation or AutoML to find a good $s$.

Response to Weakness 5:

The added results (MinCutPool, ASAP, Wit-TopoPool, MSGNN) are in global rebuttal. Our method has better classification accuracy on most of the datasets and has the highest average accuracy.

Please do not hesitate to let us know if you need more explanation or have further questions.

Dear Reviewer anD1,

We appreciate your comments and suggestions. Did our response address your concerns? We are keen to receive your feedback and provide further explanation if necessary.

Sincerely,

Authors

Response to Question 1:

Thank you for your question. We chose GIN because it is a highly representative model within the class of Graph Neural Networks (GNNs) that utilize neighbor aggregation schemes. GIN is probably the most expressive model in this category, and such schemes are widely employed in many classic GNN designs, including GCN and GraphSAGE.

In terms of theoretical analysis, our approach can be readily adapted to models like GCN and GraphSAGE with only minor modifications. GAT uses an attention mechanism for aggregation, and this approach differs significantly from the general aggregation scheme employed by GIN, GCN, and many other models. We believe the analysis of GIN provides broader and more valuable insights. Additionally, as we said in our paper, the analysis on GIN's generalization ability is currently limited, and our work aims to address this gap. It's also worth noting that our proposed reference layer can be combined with other models like GCN and GAT to classify graphs without any modification.

Response to Question 2:

Thank you for your insightful comment. While it is theoretically possible to absorb the parameter $\theta$ of the Gaussian kernel into the neural network parameters, we believe that setting or optimizing $\theta$ separately is crucial in practice.

The main reason for this is that the initial neural network parameters and the reference distributions do not necessarily ensure that nodes of graphs $\mathbf{x}$ are close to nodes of graphs $\mathbf{x}'$. If $\theta$ or the scale of the node embeddings $\mathbf{H}$ is too large, the Gaussian kernel becomes overly sharp, resulting in almost zero values. This can lead to a situation where the Maximum Mean Discrepancy (MMD) fails to effectively quantify the distance between the embeddings and the reference distributions, of which the complexity is related to $K$.

To illustrate the importance of setting or optimizing $\theta$ separately, we actually conducted experiments on the MUTAG dataset with different values of $\theta$ on Appendix D6, Table 11. The results, shown in the following table, demonstrate that the choice of $\theta$ significantly impacts classification accuracy. $\begin{matrix} \hline \mathbf{\theta} & 1\times10^{-4} & 1\times10^{-3} & 1\times10^{-2} & 1\times10^{-1} & 1 & 1\times10^{1} & 1\times10^{2} & 1\times10^{3}\\\\ \hline **Accuracy** & 0.9096 & 0.9149 & 0.9113 & 0.9113 & 0.8254 & 0.6822 & 0.5737 & 0.3345\\\\ \hline \end{matrix}$ This shows the necessity of carefully selecting or optimizing $\theta$ to ensure good classification performance. We will elaborate on these findings in the revised manuscript.

Response to Question 3:

Thank you for your insightful question. The difference arises from our model's unique use of reference distributions. In previous literature on generalization analysis, the focus has often been on models that aggregate node embeddings into vectors and then classify these vectors using another neural network, without incorporating reference distributions. This conventional approach does not typically involve a dependency on the number of classes $K$. In contrast, our model introduces reference distributions.

We have also included an additional generalization analysis theorem in the appendix (Theorem A.1) for the GIN model with mean pooling, which does not involve reference distributions. This analysis, consistent with traditional approaches, does not include $K$ in the bound. We hope this clarification and the additional theorem help highlight the unique aspects of our model and the corresponding theoretical analysis.

Please do not hesitate to let us know if you need more explanation or have further questions.

Response to Weakness 1:

Thank you for your feedback. The equation in lines 138-139 stems from the definition of the Maximum Mean Discrepancy (MMD). MMD is a measure between two discrete distributions, where the inputs are two matrices and the output is a scalar. This means that in the computation, elementary operations like summation are compulsory: it involves summation over the kernel values rather than the original node features. Regarding the claim about avoiding graph pooling operations, we would like to clarify that traditional graph pooling operations typically compress a graph's node embedding matrix into a single vector before classification. In contrast, our method allows for the direct classification of the graph using its node embeddings without compressing them into a single vector. We treat the node embedding matrix as a discrete distribution and handle it by MMD, a theoretical-grounded measure between two distributions. We appreciate your suggestion and will ensure that this distinction is more clearly communicated in the revised manuscript.

Response to Weakness 2:

Thank you for your insightful suggestion. We appreciate your interest in the impact of the discrimination loss on our model's performance. We actually included an ablation study of our proposed method without the discrimination loss in Appendix D.5 Table 10 (where the discrimination loss coefficient $\lambda = 0$). Following your second advice, we also implement a baseline (denoted as GRDL:SumDis) with discrimination loss, where node representations are first aggregated using sum-pooling and then compared with the reference distribution. The results of the ablation study regarding these two baselines are in the following table. $\begin{matrix} \hline Method & MUTAG & PROTEINS & NCI1 & IMDB-B & IMDB-M & PTC-MR & BZR & COLLAB & Average \\\\ \hline GRDL & **92.1** \pm 5.9 & **82.6** \pm 1.2 & **80.4**\pm 0.8 & **74.8**\pm 2.0 & **52.9** \pm 1.8 & **68.3** \pm 5.4 & **92.0** \pm 1.1 & **79.8**\pm 0.9 & **77.9** \\\\ GRDL:\lambda = 0 & 89.9 \pm 4.9 & 81.8 \pm 1.3 & 80.0\pm 1.6 & 73.1 \pm 1.5 & 51.3 \pm 1.4 & 66.6 \pm 5.9 & 89.5 \pm 2.3 & 79.0 \pm 1.0 & 76.4\\\\ GRDL: SumDis & 89.9 \pm 6.0 & 78.4 \pm 0.6 & 77.2 \pm 1.7 & 71.6 \pm 5.2 & 49.8 \pm 5.4 & 62.5 \pm 6.3 & 85.3 \pm 1.5 & 77.1 \pm 0.9 & 74.0\\\\ \hline \end{matrix}$

The results show that discrimination loss can increase the performance of our model. However, our method without discrimination loss still outperforms the baseline with sum-pooling and discrimination loss. We will include these results in the revised version of the manuscript.

Response to Weakness 2 (minor):

Thank you so much for pointing out these baselines. Due to the time limit, we implemented SEP, GMT, and MPNN-VN. The results are shown in the following table. Our method GRDL outperformed the competitors in most cases. We will include these additional results in our paper and discuss all the six references you mentioned. $\begin{matrix} \hline \text{Method} & MUTAG & PROTEINS & NCI1 & IMDB-B & IMDB-M & PTC-MR & BZR & COLLAB & Average \\\\ \hline \text{GRDL (ours)} & **92.1** \pm 5.9 & **82.6** \pm 1.2 & 80.4 \pm 0.8 & **74.8** \pm 2.0 & **52.9** \pm 1.8 & 68.3 \pm 5.4 & **92.0** \pm 1.1 & 79.8\pm 0.9 & **77.9** \\\\ \text{SEP [ICML 2022]} & 89.4 \pm 6.1 & 76.4 \pm 0.4 & 78.4 \pm 0.6 & 74.1 \pm 0.6 & 51.5 \pm 0.7 & 68.5 \pm 5.2 & 86.9 \pm 0.8 & **81.3** \pm 0.2 & 75.8\\\\ \text{GMT [ICLR2022]} & 89.9 \pm 4.2 & 75.1 \pm 0.6 & 79.9 \pm 0.4 & 73.5 \pm 0.8 & 50.7 \pm 0.8 & 70.2 \pm 6.2 & 85.6 \pm 0.8 & 80.7 \pm 0.5 & 75.7\\\\ \text{MPNN-VN [ICML2023]} & **92.1** \pm 5.2 & 78.3 \pm 1.0 & **80.9** \pm 0.8 & 72.4 \pm 1.2 & 50.9 \pm 1.9 & **71.4** \pm 5.2 & 90.2 \pm 1.1 & 80.1 \pm 0.8 & 77.0\\\\ \hline \end{matrix}$

Please do not hesitate to let us know if you need more explanation or have further questions.

I appreciate the detailed response which has addressed my concerns. The new ablation results (W2) further validate the proposed method and the selected comparison results seem promising. Please make sure to clarify the difference between the pre-pooling methods and your discrepancy measuring and summation strategy, and update your manuscript based on the rebuttal.

The authors referred to a PDF file in their global rebuttal, but I am unable to locate it. Could the authors please confirm whether this file has been correctly uploaded?