Gaussian Mixture Models Based Augmentation Enhances GNN Generalization

We thank Reviewer qdD6 for the feedback. In what follows, we address the raised questions and weaknesses point-by-point.

[Weaknesses] Motivation behind GMM-GDA

Indeed, other generative techniques, such as generative models, could also be used to fit the embeddings of the training data. However, we chose to use GMMs for several important reasons, particularly the efficiency and effectiveness of GMMs in this context.

GMMs are universal approximators, meaning they can effectively approximate any distribution, including the distribution of the graph embeddings. This property ensures that the augmented data is drawn from a distribution that closely aligns with the original data distribution, as shown in Theorem 3.1. While other generative methods could be used to fit the embeddings, GMMs have the advantage of being relatively simple and efficient in terms of computation (we can fit a GMM and generate new samples in very few seconds). Unlike more complex methods such as generative models, GMMs can achieve high-quality approximations with minimal computational overhead.

[Question 1] The number of Gaussian distribution

Theoretically, increasing $K$ allows the GMM to better approximate the true distribution of graph embeddings, as a higher number of components provides more flexibility in capturing complex distributions. We experimented with $K$ values in the range of 2 to 50, with the maximum $K=50$ being sufficient to balance model complexity and computational feasibility. In some cases, even a small number of components (e.g., $K= 2$ or $K=5$) was sufficient to achieve competitive performance. Below, We include a hyperparameter sensitivity analysis on the GIN backbone and the dataset IMDB--BINARY and we noticed a consistency in the results.

[Question 2] Typo

We are grateful to the reviewer for spotting the typo. We made the necessary adaptations in the submitted paper.

[Question 3] Proof of Theorem 3.1

We thank the reviewer for pointing the error in the proof. We updated the proof in the paper with the necessary changes.

[Question 4] Clarification on the definition of $\hat{\theta}_{aug}$

As defined in line 754 (in the Appendix) and in the mathematical formalism of graph data augmentation (in Section 3.1), the weights $\hat{\theta}_{aug}$ correspond to

$\text{argmin} \frac{1}{N} \sum_{n=1}^{N} \mathbb{E}_{\lambda \sim \mathcal{P} }\left [ \ell(\mathcal{G}_n^\lambda,\theta) \right ],$

and its empirically approximated by $\text{argmin} \frac{1}{N\times M} \sum_{n=1}^N \sum_{m=1}^M \ell(\mathcal{G}_n^{\lambda} {}^{n,m}, \theta).$

The proof for this formulation still holds. We have provided a clearer definition of $\hat{\theta}_{aug}$ in Section 3 of the manuscript.

[Question 5] The Upperbound in Theorem 3.1

It is true that if the maximum of the expectations $\max_{n \in \{1, \ldots,N\}} \mathbb{E}_{\lambda \sim \mathcal{P}} \left [ \left \| \mathcal{G}_n^\lambda -\mathcal{G}_n \right \| \right ], \text{ is non-zero, the Rademacher complexity of},$

$\ell_{aug}$ may not necessarily be smaller than the Rademacher complexity of $\ell$. However, by minimizing the term $\mathbb{E}_{\lambda \sim \mathcal{P}} \left [ \left \| \mathcal{G}_n^\lambda -\mathcal{G}_n \right \| \right ],$

we are more likely to reduce the additional Rademacher complexity, thus increasing the chances of achieving a lower overall Rademacher complexity. Since we use Gaussian Mixture Models (GMMs) for augmenting the graph data, which are universal approximators, we ensure that the expectations $\max_{n \in \{1, \ldots,N\}} \mathbb{E}_{\lambda \sim \mathcal{P}} \left [ \left \| \mathcal{G}_n^\lambda -\mathcal{G}_n \right \| \right ]$ approach zero.

We thank Reviewer YuEk for the feedback. In what follows, we address the raised questions and weaknesses point-by-point.

[Question 1 + Weaknesses 1] the Use of GMM Over Alternative Methods

Indeed, other generative techniques, such as generative models, could also be used to fit the embeddings of the training data. However, we chose to use GMMs for several important reasons, particularly the efficiency and effectiveness of GMMs in this context.

GMMs are universal approximators, meaning they can effectively approximate any distribution, including the distribution of the graph embeddings. This property ensures that the augmented data is drawn from a distribution that closely aligns with the original data distribution, as shown in Theorem 3.1. While other generative methods could be used to fit the embeddings, GMMs have the advantage of being relatively simple and efficient in terms of computation (we can fit a GMM and generate new samples in very few secons). Unlike more complex methods such as generative models, GMMs can achieve high-quality approximations with minimal computational overhead.

[Question 2] Distribution invariance in GMM-GDA

The decision to limit retraining to the post-readout function after augmenting the dataset is motivated by several considerations: 1/ Before the readout function, the GNN generates embeddings at the node level. Directly learning a distribution at the node level is more complex as the number of nodes varies across graphs. In contrast, the graph-level representations output are fixed-dimensional embeddings which make it easy to learn their distribution with a GMM and sample new graph representations from it. 2/ Augmenting the data increases the size of the training set by adding new graph representations, i.e. augmented graphs. However, the post-readout function is a relatively small component of the GNN, typically consisting of a linear layer followed by a softmax. As a result, training this component on the augmented dataset is computationally efficient. In contrast, retraining the message-passing layers, which involve multiple neighborhood aggregation steps and operate at the node level, would significantly increase computational training time due to the larger dataset size.

[Question 3] Analyzing the influence scores on DD Dataset

Our approach GMM-GDA, generally performs better than the baselines for the DD dataset across both GIN and GCN backbones, as demonstrated in Tables 1 and 2. Figure 2 highlights that GMM-GDA exhibits greater effectiveness, as reflected by positive influence scores, on GIN compared to GCN. This behavior is consistent not only with our method but also with the baselines, as most graph data augmentation strategies tend to enhance test accuracy more significantly for GIN than for GCN when applied on DD.

[Weaknesse 2 + Question 4] The augmentation time and training time

In addition to outperforming the baselines on most datasets, our approach offers an advantage in terms of time complexity (cf. Table 6 in Appendix E). The training time of baseline models varies depending on the augmentation strategy used, specifically, whether it involves pairs of graphs or individual graphs. Even in cases where a graph augmentation has a low computational cost for some baselines, training can still be time-consuming as multiple augmented graphs are required to achieve satisfactory test accuracy. For instance, methods like DropEdge, DropNode, and SubMix, while computationally simple, require generating multiple augmented samples at each epoch, thereby increasing the overall training time. In contrast, GMM-GDA introduces a more efficient approach by generating only one augmented graph per training instance, which is reused across all epochs. This design ensures a balance between computational efficiency and augmentation effectiveness, reducing the overall training burden while maintaining strong performance.

The only baseline that is more time-efficient than our approach is GeoMix; however, our method consistently outperforms GeoMix across all settings, as shown in Tables 1 and 2.

We thank Reviewer 25Li for the feedback. In what follows, we address the raised questions and weaknesses point-by-point.

[Weaknesses 1] On the clarity of the paper

Thank you for your feedback. We acknowledge that the clarity of Section 3 could be improved to help guide the reader more effectively through the discussion. We have added more details in the revised manuscript. We will further enhance the clarity and organization of this section in the camera-ready version of the manuscript.

[Weaknesses 2] Pre-training of the model

To clarify, our approach does not require pre-training of the model in the traditional sense. Instead, we split the training process of the GNN into two distinct parts. The GNN consists of a sequence of message-passing layers followed by a shallow neural network, referred to as the post-readout function. We have provided a detailed summary of our approach in Algorithm 1.

In our method, we first train the message-passing layers on the graph classification task. Then, we train the shallow post-readout network using the representations from both the original and augmented graphs. The second training phase, which involves the post-readout function, is very quick as it consists of a simple shallow MLP that can be trained in just a few seconds.

As a result, while the total training time is increased due to this two-step process, the additional time required for training the second part is minimal. Therefore, the overall training time is not significantly impacted by our approach.

[Weaknesses 2] Model architecture specificityl

We agree that the graph data augmentation method is dependent on the specific model architecture. The decision to perform augmentation at the level of graph representations, rather than directly on the graph inputs, is driven by Theorem 3.3. This theorem demonstrates that an effective data augmentation strategy should take into account the particular model architecture and its learned weights. Additionally, to evaluate a data augmentation strategy, it is necessary to train different GNN backbones separately on the augmented train set, e.g. evaluating a data augmentation approach on GCN and GIN requires training both GCN and GIN on the training set independently.

[Weaknesses 3] Edge insertion and Deature drop Baselines

Thank you for suggesting additional baselines. We have already included baselines [1], which outperform these methods. Nevertheless, we will add the suggested baselines in the camera-ready version of the manuscript. While edge insertion is similar to the already-used DropEdge baseline, we will add it for more complete comparaison. However, feature drop cannot be applied across all datasets, as some lack node features, making this augmentation impractical in those cases.

[1] Yoo, J., Shim, S., & Kang, U. (2022, April). Model-agnostic augmentation for accurate graph classification. In Proceedings of the ACM Web Conference 2022 (pp. 1281-1291).

[Weaknesses 4] Metric in Theorem 3.1

The norm used to measure the distance between $\mathcal{G}^\lambda$ (the augmented graph) and $\mathcal{G}$ (the original graph) can correspond to any norm defined in the space of graphs, e.g. the spectral and Frobenius norms. The specific choice of the norm should align with the Lipschitz constant $L_{\text{Lip}}$. Importantly, the result of Theorem 3.1 remains unchanged regardless of the chosen norm because all norms are equivalent in the space of graphs. If $L_{\text{Lip}}$ is taken to be the Lipschitz constant of the post-readout function only, then the norm corresponds to the difference between the graph-level embeddings produced by the readout function, i.e., $\|h_{\mathcal{G}^\lambda} - h_{\mathcal{G}}\|$.

To address this, we explicitly clarify in the revised manuscript that the norm can be general norm that match the constant of $L_{\text{Lip}}$.

[Question 1] Parameter $\lambda$ in the case of GMM

In our approach, the augmented hidden representations $h_{\mathcal{G}_{n}^{\lambda}}$ corresponds to a sampled vector from the GMM distribution $\mathcal{P}_c$ that was previously fit on the hidden representations $\mathcal{H}_c$ of the graphs in the training set with the same class $c$.

Formally, $h_{\mathcal{G}_{n}^{\lambda}}= A( \{\mathcal{H}_c\}_n, \lambda_c) = \lambda_c$, where $\lambda_c$ is sampled from the GMM distribution $\mathcal{P}_c$.

Regarding W4: The setting is clear when the distance between two graphs is defined in the embedding space, when not how it is defined? is it based on the Frobenius or spectral norm of the difference of their adjacency matrices?

Regarding Q2: Even if GMMs are universal approximators, it is not guaranteed that the samples will remain within the distribution. Firstly, the approximation capabilities depend on the number of Gaussian components used. Secondly, the EM algorithm can converge to incorrect solutions, there are even simple cases in which it gets stuck in a local minimum.

We thank Reviewer 25Li very much for their follow-up questions. In what follows, we try to further clarify the remaining questions.

[Weakness 4] Definition of Graph Distance in Non-Euclidean Spaces

Yes, the distance between two graphs can be based on the norm of the difference of their adjacency matrices, and the norm could be for example the Frobenius or spectral norm. As mentioned in our response, the inequality holds for both norms since all norms are equivalent in finite-dimensional spaces. Specifically, let us consider the graph space $(\mathbb{G}, \lVert \cdot \rVert_{\mathbb{G}})$ and the feature space $(\mathbb{X}, \lVert \cdot \rVert_{\mathbb{X}})$, where $\lVert \cdot \rVert_{\mathbb{G}}$ and $\lVert \cdot \rVert_{\mathbb{X}}$ denote the norms applied to the graph structure and features, respectively. Assuming a maximum number of nodes per graph, which is a realistic assumption for real-world data, the product space $\mathbb{G} \times \mathbb{X}$ is a finite-dimensional real vector space, and all the norms are equivalent. Thus, the choice of norm does not affect the theorem, as long as the Lipschitz constant is adjusted accordingly.

When considering only structural changes, with fixed node features, the distance between two graphs $\mathcal{G}^\lambda,\mathcal{G}$ is defined as: $\lVert\mathcal{G}^\lambda - \mathcal{G} \rVert= \lVert A - A^\lambda \rVert_{\mathbb{G}}, \qquad \qquad (1)$ where $A,A^\lambda$ are respectively the adjacency matrix of $\mathcal{G}^\lambda,\mathcal{G},$ and the norm $\lVert\cdot \rVert_{\mathbb{G}}$ can be the Frobenius or spectral norm. If both structural and feature changes are considered, the distance extends to: $\lVert\mathcal{G}^\lambda - \mathcal{G} \rVert= \alpha \lVert A - A^\lambda \rVert_{\mathbb{G}} +\beta \lVert X - X^\lambda \rVert_{\mathbb{X}}, \qquad \qquad (2)$ where $X^\lambda, X$ are the node feature matrices of $\mathcal{G}^\lambda,\mathcal{G}$ respectively, and $\alpha, \beta$ are hyperparameters controlling the contribution of structural and feature differences, respectively.

In most baselines graph augmentation techniques, such as for instance $\mathcal{G}$-Mixup, SubMix, and DropNode, the alignment between nodes in the original graph $\mathcal{G}$ and the augmented graph $\mathcal{G}^\lambda$ is known. However, in cases where the node alignment is unknown, we must take into account the node permutations. The distance between the two graphs is then defined as:

$

\lVert\mathcal{G}^\lambda - \mathcal{G} \rVert= \min_{P \in \Pi} \left( \alpha \lVert A - P A^\lambda P^T \rVert_{\mathbb{G}} + \beta \lVert X - P X^\lambda \rVert_{\mathbb{X}} \right), \qquad \qquad (3)

$

where $\Pi$ is the set of permutation matrices. The matrix $P$ corresponds to a permutation matrix used to order nodes from different graphs. By using Optimal Transport, we find the minimum distance over the set of permutation matrices, which corresponds to the optimal matching between nodes in the two graphs. This formulation represents the general case of graph distance, which has been used in the literature [1].

Importantly, Theorem 3.1 applies to both Frobenius and spectral norms in the three scenarios (1), (2) and (3) that we lay out in our answer. Hence, we tried to state the theorem in full generality in our manuscript. But we appreciate that the current statement may be perceived to lack specificity. We have provided more details and insights on the norms for which Theorem 3.1 applies to in the revised version of the manuscript (in Section 3.1 Lines 229-35 and in Appendix G).

[1] Abbahaddou, Y., Ennadir, S., Lutzeyer, J. F., Vazirgiannis, M., & Boström, H. (2024)."Bounding the Expected Robustness of Graph Neural Networks Subject to Node Feature Attacks." In The Twelfth International Conference on Learning Representations.

Table 1: Ablation Study GIN

Table 2: Ablation Study GCN

We compare these approaches for both GCN and GIN in the table above. As noticed, GMM with EM consistently outperforms the alternative methods across most datasets in terms of accuracy. The VBI method, an alternative approach for estimating GMM parameters, yields comparable performance to the EM algorithm. This consistency across datasets highlights the effectiveness and robustness of GMMs in capturing the underlying data distribution.

In certain cases, particularly with the GIN model, we observed competitive performance from the GAN approach, which, unlike GMM, requires additional training. Hence, GMMs provide a more straightforward and efficient solution.

We have included this detailed ablation study in the new version of the manuscript, c.f. Appendix D.

[2] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016.

[3] Guo, H., Lu, C., Bao, F., Pang, T., Yan, S., Du, C., |& Li, C. (2024). Gaussian mixture solvers for diffusion models. Advances in Neural Information Processing Systems, 36.

[4] Härdle, W., Werwatz, A., Müller, M., Sperlich, S., Härdle, W., Werwatz, A., et al. & Sperlich, S. (2004). Nonparametric density estimation. Nonparametric and semiparametric models, 39-83.

[5] Roger B. Nelsen (1999), "An Introduction to Copulas", Springer. ISBN 978-0-387-98623-4

[6] Dimitris G Tzikas, Aristidis C Likas, and Nikolaos P Galatsanos. (2008).The variational approximation for bayesian inference. IEEE Signal Processing Magazine, 25(6):131–146,

[7] Miin-Shen Yang, Chien-Yo Lai, and Chih-Ying Lin. A robust em clustering algorithm for gaussian mixture models. Pattern Recognition, 45(11):3950–3961, 2012.

[8] Xu, L., & Veeramachaneni, K. (2018). Synthesizing tabular data using generative adversarial networks. arXiv preprint arXiv:1811.11264.

[Question 2] Consistency with the Number of Gaussian Components

We experimented with values of $K$ in the range from 2 to 50, finding that $K=50$ was sufficient to balance model complexity and computational feasibility. In many cases, even a small number of components (e.g., $K=2$ or $K=5$) yielded competitive performance (see Table 9 in the submitted manuscript). Below, we include a hyperparameter sensitivity analysis conducted on the GIN backbone using the IMDB-BINARY dataset, where we observed consistent results.

We thank Reviewer 7B8Ku for the feedback. In what follows, we address the raised questions and weaknesses point-by-point.

[Weaknesses 1] Graph-Specificity of our Approach

We acknowledge that our proposed data augmentation method is not inherently specific to graphs and could indeed be applied to other data types. However, we view this as a significant strength. Working on the field of Graph Machine Learning, we chose to apply this method to Graph Neural Networks (GNNs) because of its potential to enhance the performance of graph-based models. Moreover, several established graph augmentation methods, such as SubMix, GeoMix, and GraphMix, are adaptations of general data augmentation strategies originally designed for other domains, such as Mixup. These approaches have proven to be effective in graph contexts despite their general origins.

Notably, the choice of augmenting the train dataset at the level of graph representation, and not at the level of graph input, is motivated by Theorem 3.3, which shows that a data augmentation strategy should depend on the specific model architecture and its weights.

Another key motivation is that, unlike many existing baselines which focus on augmenting either the graph structure (e.g., perturbing edges) or the node features separately, our method operates directly on graph representations. This enables us to simultaneously augment both structural and feature information, leading to improved generalization in GNNs.

[Weaknesses 2] Improvements of our Approach

We trained all baseline models using the same train/validation/test splits, GNN architectures, and hyperparameters to ensure a fair comparison. It is worth noting that the baselines also exhibit high standard deviations, which is a common characteristic in graph classification tasks. Unlike node classification, graph classification is known to have larger variance in performance metrics [1].

In addition to outperforming the baselines on most datasets, our approach offers an advantage in terms of time complexity (cf. Table 6 in Appendix E). The training time of baseline models varies depending on the augmentation strategy used, specifically, whether it involves pairs of graphs or individual graphs. Even in cases where a graph augmentation has a low computational cost for some baselines, training can still be time-consuming as multiple augmented graphs are required to achieve satisfactory test accuracy. For instance, methods like DropEdge, DropNode, and SubMix, while computationally simple, require generating multiple augmented samples at each epoch, thereby increasing the overall training time. In contrast, GMM-GDA introduces a more efficient approach by generating only one augmented graph per training instance, which is reused across all epochs. This design ensures a balance between computational efficiency and augmentation effectiveness, reducing the overall training burden while maintaining strong performance.

[1] Bianchi, F. M., Lachi, V. (2024). The expressive power of pooling in graph neural networks. Advances in neural information processing systems, 36.

[Question 1] Explanation of Theorem 3.1.

In Theorem 3.1, we established a mathematical framework to connect graph data augmentation with its impact on the generalization of GNNs. To evaluate GNN generalization, we employed the concept of Rademacher Complexity to derive a regret bound on the generalization gap. Rademacher Complexity, a fundamental concept in statistical learning theory, quantifies a model's capacity to generalize to unseen data by assessing its ability to fit random labels or noise. Better generalization is associated with a lower Rademacher Complexity. Theorem 3.1 demonstrates that data augmentation can reduce Rademacher Complexity, i.e., $\mathcal{R}(\ell_{aug}) \leq \mathcal{R}(\ell)$. Specifically, the augmented graph should achieve a lower upper bound value for the following expression: $\max_{n \in \{1, \ldots,N\}} \mathbb{E}_{\lambda \sim \mathcal{P}} \left [ \left \| \mathcal{G}_n^\lambda -\mathcal{G}_n \right \| \right ]$ which implies that the augmented data must follow the same distribution as the original graph representation. This condition is ensured by Gaussian Mixture Models (GMMs), which are universal approximators.

Given the brevity of your review, we hope that our additional explanations have helped you understand our work more deeply and give you the opportunity to reevaluate your assessment of our work.