Graph Neural Network Generalization With Gaussian Mixture Model Based Augmentation
Authors: Yassine ABBAHADDOU
Fragkiskos D. Malliaros
Johannes F. Lutzeyer
Amine M. Aboussalah
Michalis Vazirgiannis
Venue: ICML 2025 poster (2025)
Keywords: Graph Neural Networks
Generalization
Data Augmentation
Graph Classification
Rademacher complexity
Primary area: deep_learning->graph_neural_networks
TL;DR: We introduce GRATIN, a new efficient graph data augmentation scheme based on GMMs to improve Graph Neural Network generalization.

Abstract: Graph Neural Networks (GNNs) have shown great promise in tasks like node and graph classification, but they often struggle to generalize, particularly to unseen or out-of-distribution (OOD) data. These challenges are exacerbated when training data is limited in size or diversity. To address these issues, we introduce a theoretical framework using Rademacher complexity to compute a regret bound on the generalization error and then characterize the effect of data augmentation. This framework informs the design of GRATIN, an efficient graph data augmentation algorithm leveraging the capability of Gaussian Mixture Models (GMMs) to approximate any distribution. Our approach not only outperforms existing augmentation techniques in terms of generalization but also offers improved time complexity, making it highly suitable for real-world applications.

Graph Neural Network Generalization with Gaussian Mixture Model Based Augmentation
Yassine Abbahaddou' 1
Fragkiskos D. Malliaros Johannes F. Lutzeyer
Amine M. Aboussalah Michalis Vazirgiannis 14
Abstract
data, GNNs have been developed as specialized class of neural networks designed to operate directly on graphs. Unlike traditional neural networks that are optimized for grid-like data, such as images or sequences, GNNs are engi- neered process and learn from the relational information embedded in graph structures. GNNs have demonstrated latt-off-hhe-art performance across range of graph rep- resentation learning tasks such as node and graph classi- fication, proving their effectiveness in various real-world applications (Vignac et al., 2023; Corso et al., 2023; Duval et al., 2023; Castro-Correa et al., 2024; Chi et al., 2022; Panagopoulos et al., 2024; Aboussalah & Ed-dib, 2025).
Graph Neural Networks (GNNs) have shown great promise in tasks like node and graph classi- fication, but they often struggle to generalize, par- ticularly to unseen or out-of-distribution (OOD) data. These challenges are exacerbated when training data is limited in size or diversity. To address these issues, we introduce a theoretical framework using Rademacher complexity to com- pute a regret bound on the generalization error and then characterize the effect of data augmen- tation. This framework informs the design of GRATIN, efficient graph data augmentation algorithm leveraging the capability of Gaussian Mixture Models (GMMs) to approximate any dis- tribution. Our approach not only outperforms existing augmentation techniques in terms of gen- cralization but also offers improved time com- plexity, making it highly suitable for real-world applications. Our code is publicly available at: http://gtthh..commmbaaaaddddddt
Despite their impressive capabilities, GNNs face significant challenges related to generalization, particularly when han- dling unseen out-of-distribution (OOD) data (Guo et al., 2024; Li et al., 2022). OOD graphs are those that differ significantly from the training data in terms of graph struc- ture, node features, or edge types, making difficult for GNNs adapt and perform well on such data. This chal- lenge also faced when GNNs are trained on small datasets, where the limited data diversity hampers the model's 's ability to generalize effectively. To address these challenges, the community has explored various strategies to improve the robustness and generalization ability of GNNs (Abbahaddou et al., 2024; Yang et al., 2023).
1. Introduction
Graphs are a fundamental and ubiquitous structure for mod- eling complex relationships and interactions. In biology, graphs are employed to represent complex networks of pro- tein interactions and in drug discovery by modeling molecu- lar relationships (Gaudelet et al., 2021; Jagtap et al., 2022). Similarly, social networks capture relationships and commu- nity interactions (Aboussalah et al., 2023a; Zeng et al., 2022; Malliaros & Vazirgiannis, 2013; Newman al., 2002). To address the unique challenges posed by graph-structured
Generalization bounds for GNNs have been derived using various theoretical tools, such the Vapnik-Chervonenkis (VC) dimension (Pfaff et al., 2021; Garg et al., 2020) and Rademacher complexity (Yin et al., 2019; Esser et al., 2021). Furthermore, Liao et al. (2021) were among the first to establish generalization bounds for GNNs using the PAC- Bayesian approach. Neural Tangent Kernels have also been employed to study the generalization properties of infinitely wide GNNs trained via gradient descent (Jacot et al., 2018 8; Du et al., 2019; Huang et al., 2024). While most existing research has focused on the node classification task, enhanc- ing generalization in graph classification presents unique challenges. Techniques to improve generalization in graph classification can broadly categorized into architectural and dataset-based strategies (Tang & Liu, 2023; Buffelli et al., 2022). On the dataset side, techniques like adversar-
"Work partially completed during research visit to NYU. LLIX École Polytechnique. IP Paris Université Paris-Saclay, CentraleSupélle, Inria 'NYU Tandon School of Engineering "MBZUAI. Correspondence to: Yassine Abbahaddou <yas- sinc.abbahaddoue @ polytchniquee.duuu..
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
ial training and data augmentation play significant role. More broadly, data augmentation methods create synthetic or modified graph instances to enrich the training set, reduc- ing overfitting and enhancing the model's adaptability to diverse graph structures. Data augmentation has shown its benefits across different types of data structures such as im- ages (Krizhevsky et al., 2012) and time series (Aboussalah et al., 2023b). For graph data structures, generating aug- mented versions of the original graphs, such as by adding or removing nodes and edges or perturbing node features (Rong et al., 2020; You et al., 2020), allows for the creation of a more varied training set. Inspired by the success of the Mixup technique in computer vision (Rebuffi et al., 2021; Dabouei et al., 2021; Hong et al., 2021), additional meth- ods such as G-Mixup and GeoMix have been developed to adapt Mixup for graph data (Ling et al., 2023; Han al., 2022). These techniques combine different graphs to cre- ate new, synthetic training examples, further enriching the dataset and enhancing the GNN's ability to generalize to new unseen graph structures.
2. Background and Related Work
Notation. Let G = (V,E) denote graph, where V repre- sents the set of vertices and & represents the set of edges. We use p | to denote the number of nodes. For a node v E V, let (v) be the set of its neighbors, defined as N(t) = {u: (v,u) E E). The degree of vertex is the number of neighbors it has, which is deg(v) = [.()) A graph is commonly represented by its adjacency matrix A E RPXP, where the (i.j)-th element of this matrix is equal to the weight of the edge between the i-th and j-th node of the graph and weight of zero in case the edge does not exist. Additionally, in some cases, nodes may have associated feature vectors. We denote these node features by X E Rexd where d is the dimensionality of the features.
Graph Neural Networks (GNNs). A GNN model con- sists of multiple neighborhood aggregation layers that use the graph structure and the feature vectors from the previ- ous layer generate updated representations for the nodes. Specifically, GNNs update a node's 's feature vector by aggre- gating information from its local neighborhood. Consider a GNN model with T neighborhood aggregation layers. Let he (0) denote the initial feature vector of node v, which is the corresponding row in X. At each layer t > O, the hidden state h of node e is updated as follows:
In this work, we introduce GRATIN, a graph augmenta- tion technique based on Gaussian Mixture Models (GMMs), which operates the level of the final hidden representa- tions. Specifically, guided by our theoretical results, we ap- ply the Expectation- Expeetation-Maximizztton (EM) algorithm to train a GMM on the graph representations. We then use this GMM to generate new augmented graph representations through sampling, enhancing the diversity of the training data.
m AGGREGATE()) (t) (fn u uE N()))) h COMBINE()) (hC-", m!)).
The contributions of our work are as follows:
where AGGREGATE() is a permutation-invarite function that combines the feature vectors of t's neighbors into an aggregated vector. This aggregated vector, together with the previous feature vector h!--1, is fed to the COMBINE({ function, which merges these two vectors to produce the updated feature vector of v.
Theoretical framework for generalization in GNNs. We introduce a theoretical framework that allows us to rigorously analyze how graph data augmentation im- pacts the generalization GNNs. This framework offers new insights into the underlying mechanisms that drive performance improvements through augmentation.
After T iterations of neighborhood aggregation, GNN typically produces a graph-level representation by first ap- plying a permutation-invaaiin operator, to the final node embeddings. This aggregated output subsequently passed through a trainable neural network, such as a multi-layer perceptron (MLP), denoted by Vi and referred to as the post-readoul neural network, to produce the final graph-level predictions or representations. Mathematically, this process can be expressed as
Efficient graph data augmentation via GMMs. We propose GRATIN, a fast and efficient graph data aug- mentation technique leveraging GMMs. This approach enhances the diversity of training data while maintain- ing computational simplicity, making it scalable for large graph datasets.
readout function, e.g., Sum
Comprehensive theoretical analysis using influence functions. We perform an in-depth theoretical analysis of our augmentation strategy through the lens of influ- ence functions, providing a principled understanding of the approach's 's impact generalization performance.
hg = V o READOUT ((n!.) : VEV)).
Two popular GNN architectures
Graph Convolution Networks (GCN) and Graph Isomorphism Networks (GIN) (Kipf & Welling, 2017; Xu et al., 2019). The exact expres- sion of these models can be found in Appendix C.
Empirical Validation. Through experiments on real- world datasets we confirm GRATIN to be a fast, high- performing graph augmentation scheme in practice.
Data Augmentation for Graphs. Graph data augmentation has become essential to enhance the performance and robust-
2

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
ness of GNNs. Classical graph augmentation techniques focus structural modifications to generate augmented graphs. Key methods here include DropEdge, DropNode, and Subgraph sampling (Rong et al., 2020; You et al., 2020). For instance, DropEdge randomly removes a subset of edges from the graph during training, improving the model's 's ro- bustness to missing or noisy connections. Similarly, DropN- ode removes certain nodes as well as their connections, assuming that the missing part of nodes will not affect the semantic meaning, i.e., the structural and relational informa- tion of the original graph. Subgraph sampling, on the other hand, samples a subgraph from the original graph using random walks to use as a training graph.
3.1. Formalism of Graph Data Augmentation
We focus on the task of graph classification, where the objective is to classify graphs into predefined categories. Let D denote the distribution of graphs. Given training set of graphs Dirann {( Vn)) n = 1..... .N) Gn is the n-th graph and yn is its corresponding label belonging to a set {O...., CO.....C. Each graph Gn is represented as a tuple (Va, En, Xn). where Vn, denotes the set of nodes with cardinality Pn [ial. Ex V,, Vn is the set of edges, and X,, E RP- is the node feature matrix of dimension d. The objective is to train GNN f(- 0) that can accurately predict the class labels for unseen graphs in the test set Dest = = n = 1,..., -Matt. The classical training approach involves minimizing the following loss function,
Beyond classical methods, recent advancements have ex- plored more sophisticated augmentation techniques, focus- ing on manipulating graph embeddings and leveraging the geometric properties of graphs. Following the effectiveness of the Mixup technique in computer vision (Rebuffi et al., 2021; Dabouei et al., 2021; Hong et al., 2021), several works describe variations of the Mixup for graphs. For example, the Manifold-Mixup Mixup model conducts a Mixup operation for graph classification in the embedding space. This tech- nique interpolates between graph-level embeddings after the READOUT function, blending different graphs in the embedding space (Wang et al., 2021). Similarly, G-Mixup (Han et al., 2022) uses graphons to model the topologi- cal structures of each graph class and then interpolates the graphons of different classes, subsequently generating syn- thetic graphs by sampling from mixed graphons across dif- ferent classes. It is important note that G-Mixup operates under a significant assumption: graphs belonging to the same class can be produced by single graphon. Other advanced techniques include S-Mixup method, which in- terpolates graphs by first determining node-level correspon- dences between a pair of graphs (Ling et al., 2023), and FGW-Mixup, which adopts the Fused Gromov- Wasserrtenn barycenter to compute mixup graphs but suffers from heavy computation time (Ma et al., 2024). Finally, GeoMix (Zeng et al., 2024) leverages Gromov- Wasserstein geodesics to interpolate graphs more efficiently. By leveraging these structural augmentation techniques, GNNs can better gener- alize to unseen graph structures.
N
0),
where E denotes the cross-entropy loss function.
To improve the generalization performance of GNNs, we introduce a graph data augmentation strategy. For each training graph G. in the dataset, we generate M augmented graphs, denoted as {9n.. Vn.m | m = M.. }. where M is the number of augmented graphs generated per train- ing graph. These augmented graphs are obtained using a graph augmentation generator Ax, parameterized by À as a mapping A Gns Vo. + Ax (Gn, Bo.) E G x Y, where G denotes the space of all possible graphs, and Y is the label space. The generator A may be either deterministic or stochastic with a dependence on a prior distribution P(A). Examples of such augmentation strategies can be found in Appendix G.
We use the notation TThh AA to represent augmented graph sampled from the augmentation strategy Ax. i.e., ( """") ~ (x yn). With the augmented data, the loss function is modified to account for multiple augmented versions of each graph,
rang
Σ B BornA [ 0). Th n=1
For simplicity, we denote the loss for the original graph by (( U), = eg 0) and the loss for an augmented graph as,
3. GRATIN: Gaussian Mixture Model for Graph Data Augmentation
E E Ax [UUis (U((T,0)),0 (T)))) ( 0). Via the law of large numbers, L is empirically estimated,
In this section, we introduce the mathematical framework for graph data augmentation and its connection to the gener- alization of GNNs. Then, we present our proposed model GRATIN, which is based on GMMs for graph augmentation.
Lang
gaue (G., 0)
N M NM ΣΣ m=11
E(f(ā P).
3

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
To understand the impact of data augmentation on the graph classification performance, we analyze the effect of the aug- mentation strategy on the generalization risk Bg-D [e(g, O)]. More specifically, we want to study the generalization error,
that the input node features and graph structures in real- world datasets are typically bounded, i.e., node features are typically normalized or constrained within fixed range, while graph structures, represented by adjacency matrices or their normalized forms, have bounded spectral properties, ensuring a constrained input space. Additionally, we can ensure that the loss function is bounded within [o. 1] by composing any standard classification loss with a strictly increasing function that maps values to the interval [0, 1]. We provide the proof of this theorem in Appendix A.
17 = Eg [e(g, amgg) - -g-D (e(g, 8.)],
where mmg and 0. are the optimal GNN parameters for the augmented and con-augmented settings, respectively,
0. arg min Ce, Dang arg min Cy
6
A direct implication of Theorem 3.1 is that if we chose the right data augmentation strategy AĄ that minimizes the expected distance between original graphs and augmented ones EG-DD ||| g we can guarantee with high probability that the data augmentation decreases both the Rademacher complexity and the generalization risk. Specif- ically, we need enough diversity to ensure a non-zero dif- ference in the Rademacher complexity compared to the non-augmented case while controlling the expected distance that we have a chance reduce both the Rademacher complexity and the overall generalization error 1). On the other hand, if the distance is large, we cannot guarantee that data augmentation will outperform the normal training setting.
which can be estimated empirically as follows,
N
i- arg min
((Gn.
0
M Σ 72 m=
N
0ang = arg min
M
{ 0).
NM
By theoretically studying the generalization enror 77. we aim to quantify the effect of each augmentation strategy on the overall classification performance, providing insights into the benefits and potential trade-offs of data augmentation in graph-based based learning tasks. In Theorem 3.1, we present a regret bound of the generalization error using Rademacher complexity defined as follows (Yin et al., 2019),
The findings of Theorem 3.1 hold for all norms defined on the graph input space. Specifically, let us consider the graph structure space (A, I1-I1a) and the feature space (X, Il-lIx), where II-IIA and Il-lIx denote the norms applied to the graph structure and features, respectively. Assuming a maximum number of nodes per graph, which is realistic assumption for real-world data, the product space A x X is a finite- dimensional real vector space, and all the norms are equiva- lent. Thus, the choice of norm does not affect the theorem, as long as the Lipschitz constant is adjusted accordingly. Additional details on graph distance metrics and the com- parison between the Lipschitz constants of GCN and GIN are provided in Appendices H and C.
N
R()) Be,,
sup AE
Σ CGE(Gn N
1
where 6, are independent Rademacher variables, taking val- ues +1 or - with equal probability, P. is the Rademacher distribution, and è is the hypothesis class. Rademacher complexity is fundamental concept in statistical learning, which indicates how well learned function will perform on unseen data (Shalev-Shwartz & Ben-David, 2014). Intu- itively, the Rademacher complexity measures the capacity of a GNN to fit random noise (Zhu et al., 2009), where a lower Rademacher complexity indicates better generalization.
The variety of distance metrics on the original graph space offers different upper bounds, which leads to distinct crite- ria for data augmentation based on these metrics to control the upper bound in Theorem 3.1. Instead, we propose to shift the focus to the hidden representation space of graphs, where we aim to derive more consistent and meaningful augmentations. If Ltip is taken to be the Lipschitz constant of the post-readout function only, then Theorem 3.1 still ap- plies for the norm on the graph-level embeddings produced by the readout function, i.e., llb hg ||-
Theorem 3.1. Let be a classification loss function with Laip as a Lipschitz constant and ((.,)) E [0, 1]. Then, with a probability at least 8 over the samples Drrain we have,
E(G [ - EgND 0() < 2R(4 + EgND
2 llg((4/))
|| --
///
+ 2p B AwA
N
Moreover, we have,
R((aau)) S R(1)= + max Emm || - 0.
Working at the level of hidden representations of graphs, rather than directly in the graph input space, offers addi- tional advantages. Hidden representations capture both the structural information and node features of each graph, en-
Theorem 3.1 relies on the assumption that the loss function Lipschitz continuous. This assumption is realistic, given
4

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
abling augmentation that enhances the generalization of both aspects simultaneously. Moreover, node alignment is needed to compare the original and augmented graphs, which is computationally expensive. By operating on hidden representations instead, node alignment becomes unneces- sary. Furthermore, as we will discuss in Section 3.4, the ef- fectiveness of augmented data depends on the specific GNN architecture. By leveraging graph representations learned through GNN, we ensure that the augmentation process remains architecture-specffe aligning with the inductive biases of the chosen model.
Algorithm 1 Graph classification with GRATIN
Inputs: GNN of T layers f(::0) = READCUT o g. where is the composition of message passing layers, i.e, g uu (AGGREGATE o COMBINE("(( and v is the post- readout function, graph classification dataset D, loss function L: Steps:
1. Train GNN f on the training set Dunin
2. Use the trained message passing layers and the readout function
to generate graph representations H = (he,, s.t. Gu. E Daminn for the training set;
3. Partition the training set Dimin by classes, such that Dunin
U. De where De (G Daaiin y/ = c]:
foreach E {O...., C} do
3.1. Fit a GMM distribution Pa on the graph representations
3.2. Proposed Approach
no= {hg, s.t. Gn E DD);
3.2. Sample new graph representation H. = (hi s.t. h ~ pe}
Based on the theoretical findings, it is crucial to employ a data augmentation technique that effectively controls the term L...D.....A [lbb- - hell which measures the ex- pected deviation in the hidden representations of graphs under augmentation, to achieve stronger generalization guar- antees. To better understand this term, we express it in terms of graph representations rather than graphs themselves,
from the distribution ps:
3.3. Include the sampled representations He with trained representations H. = H. U Hai end foreach
4. Finetune the post-readout function Ti on the graph classification task directly on the new training set H = UcHei
hill where Qx represents the augmentation strategy at the level of hidden representations, replacing Ax, which operates at the graph level, and óp h + * > Sne, (h) represents the Dirac distribution over the training graph representations, capturing the empirical distribution of graph embeddings.
Bo-DDaA [llb- - hell = B..... R
These embeddings are used as the basis for generating augmented training graphs. We then partition the train- ing set Duain by classes, such that D U. D. where D. = (G,, E Duriin Yn c). The objective is to learn new graph representations from these embeddings, and create augmented data for improved training.
Various universal approximators can be used generators Qx. including generative models like Generative Adversar- ial Networks (GANs) (Yang et al., 2012). These models are capable of approximating any probability distribution, making them powerful tools for learning complex augmen- tations. However, we specifically choose Gaussian Mixture Models (GMMs), which are well-suited for this purpose, and can effectively approximate any data distribution, c.f. Theorem 3.2. GMMs are computationally fast compared to other generative approaches, making them suitable for sarge-scale graph datasets. As shown in Appendix E, GMM- based augmentation yields better results compared to alter- native generative strategies. Moreover, due to the exponen- tial decay of Gaussian distributions the expected distance E B inQqx [|h- hi]] is well-controlled.
We use the EM algorithm to learn the best-fitting GMM for the embeddings of each partition De, denoted as He (hg,, s.t. 9, E D.). The EM algorithm finds maximum likelihood estimates for each cluster He, following the pro- cedure described in (Bishop & Nasrabadi, 2006).
Once a GMM distribution Pe is fitted for each partition Dec we use this GMM to generate new augmented data by sampling hidden representations from Pe. Each new sample drawn from Pe then assigned the corresponding partition label c, ensuring that the augmented data inherits the label structure from the original partitions. After merging the hidden representations of both the original training data and the augmented graph data, we fine-tune the post-readout function, i.e., the final part of the GNN, which occurs after the readout function, on the graph classification task. Since the post-readout function consists of a linear layer followed by a Softmax function, the finetuning process is relatively fast. To evaluate our model during inference on test graphs, we input the test graphs into the GNN layers trained in the initial step to compute the hidden graph representations. For the post-readout function, we use the weights obtained from the second stage of training. Algorithm I and Figure I provide a summary of the GRATIN model.
Theorem 3.2. (Goodfellow et al., 2016, Page 65) A Gaus- sian mixture model a universal approximator of densities, in the sense that any smooth density can be approximated with any specific nonzero amount of error by Gaussian mixture model with enough components.
To achieve this, we first train a standard GNN on the graph classification task using the training set. Next, we obtain embeddings for all training graphs using the READOUT output, resulting in H = (hg., s.t. 9n Duain).
5

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
Graph Representations
Nodes
Edges
Trained Representations
Sampled Augmented
Node
Representations
AGGREGATE
COMBINE
READOUT
Of
REDOUT
Post-Readout
Sample
Fit
Training of Message Passing Layers
p(h)
NyN(h
(4k, Ex)
Set of train graphs from the same class
Figure 1. Illustration of GRATIN. Step 1. We first train the GNN on the graph classification task using the training graphs. Step 2. Next, we utilize the weights from the message passing layers generate graph representations for the training graphs. Step 3. A GMM then fit to these graph representations, from which we sample new graph representations. Step 4. Finally, we fine-tune the post-readout function for the graph classification task, using both the original training graphs and the augmented graph representations. For inference on the test set, we use the message passing weights trained in Step I and the post-readout function weights trained in Step 4.
3.3. Time Complexity
3.4. Analyzing the Generalization Ability of the Augmented Graphs via Influence Functions
One advantage of GRATIN is its efficiency, it gener- ates new augmented graph representations with minimal computational time. Unlike baseline methods, which apply augmentation strategies to each individual training graph (or pair of graphs in Mixup-basee approaches) separately, our method learns the distribution of graph representations across the entire training dataset simultaneously using the EM algorithm (Ng, 2000). If N [Prain| is the num- ber of training graphs in the dataset, d is the dimension of graph hidden representations (hg. G E Duain, and K is the number of Gaussian components in the GMM, then the complexity to fit a GMM on iterations is O(N K T -f) (Yang et al., 2012). We compare the data augmentation times of our approach and the baselines in Table 7. Due to our different training scheme, i.e., where we first train the message passing layers and then train the post-readout function after learning the GMM distribution, we have mea- sured the total backpropagation time and compared it with the backpropagation time of the baseline methods. The training time of baseline models varies depending on the augmentation strategy used, specifically whether it involves pairs of graphs individual graphs. Even in cases where graph augmentation has low computational cost for some baselines, training can still be time-consuming as multiple augmented graphs are required to achieve satisfactory test accuracy. In contrast, GRATIN generates only one aug- mented graph per training graph, demonstrating effective generalization on the test set. Overall, our data augmen- tation approach is highly efficient during the sampling of augmented data, with minimal impact on the training time. A complete analysis of the time complexity of GRATIN and the baselines can be found in Appendix F.
We use influence functions (Law, 1986; Koh & Liang, 2017; Kong et al., 2021) to understand the impact of augmented data on the model performance on the test set and thus motivate the use of a data augmentation strategy, which is specific to the model architecture and model weights.
In Theorem 3.3, we derive closed-form formula for the impact of adding an augmented graph G.I the GNN's performance on a test graph ghst where the GNN is trained solely on the training set, without the augmented graph.
Theorem 3.3. Given test graph Gk, let 0 = arg ming C be the GNN parameters that minimize the objective function in (1). The impact of upweighting the objective function to Lane = L + nn (( 0), where 9Thh is an augmented graph candidate of the training graph In and Cn.n a sufficiently small perturbation parameter, on the model per- formance the test graph G is given by
de(g AE(G)e, VollGk ,0H, Voll" 0), dmn.
where den. arg ming Laug denotes the parameters that minimize the upweighted objective function Lavs, and H V2,(()) is the Hessian matrix of the loss w.r.l. the model parameters.
We provide the proof of Theorem 3.3 in Appendix B. The influence scores useful for evaluating the effectiveness of the augmented data on each test graph. The strength of influence function theory lies in its ability to analyze the effect of adding augmented data to the training set without actually retraining on this data. As noticed, these influence scores depend not only on the augmented graphs themselves
6

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
but also on the model's weights and architecture. This highlights the need for graph data augmentation strategy tailored specifically to the GNN backbone in use, opposed to traditional techniques like DropNode, DropEdge, and G-Mixup, which are general-purpose methods that can be applied with any GNN architecture.
Table I. Classification accuracy (+ std) on different henchmark graph classification datasets for the data augmentation baselines based on the GCN backbone. The higher the accuracy (in %) the
better the model. Highlighted are the first, second best results.
Model
IMDB-BIN
IMDB-MUL
MUTAG
PROTEINS
DD
No Aug.
23.00
47.73:224
73.9215.00
69.99
69.69:2.49
Dropkige
71.701.540 e0
45.67.22.4
73.39 s6
20.07
69. 69:35,0.a7
Theorem 3.3 is valid for any differentiable loss function. More specifically, if the chosen loss is the cross entropy or the negative log-likelihoon then the Hessian matrix cor- responds to the Fisher information matrix (Barshan et al., 2020; Lee et al., 2022). Consequently, the norm of H i.., the inverse of the Hessian matrix, can be bounded above using the Cramér- Rao inequality (Nielsen, 2013). There- fore, a trivial case where the norm of influence scores is zero arises when the gradient of the loss function with respect to the input graphs vanishes. This scenario, for instance, can occur in the DD dataset when using GIN. A detailed analysis of this phenomenon is provided in Section 4. In these cases, data augmentation becomes ineffective, having minimal impact on the GNN's ability to generalize. We can measure the average influence I(( of an augmented graph S.te on the test set by averaging the derivatives as follows,
DropNode
74
43.80199.4
73.89aa.ss
=4.81=400
69 69:01=8.9
SuhMix
72.70.3.00
46.00...44
77.13...m
67.57 1:5
70.11 1s
G-Mixup
72.10+a.00
48 18.33-9.06
88. 7
65 68.66888.00
20
GeoMix
09.69
19. 19.80=4771
71.39
69. 63
68. 50
GrATIN
71.00s4.0
49.82.4.20
76.06
70.97ss.07
71 91.90:9.11
Table 2. Classification accuracy (#: std) on different benchmark graph classification datasets for the data augmentation baselines based on the GIN backbone. The higher the accuracy (in %) the better the model. Highlighted are the first, second best results.
Model
IMDB-BIN
IMDB-MUL
MUTAG
PROTHINS
DD
No Aug.
70.30 00
48 05
13.4222.19
08.54
68 00:20 is
DropEdge
70. 10=4.0
16.80=539
74.88=a.52
68 27
67 82
DrupNode
20.30=1.41
45. 25.20:424
75.50e7au
45.40:47:
69 89.01:3.00
SubMix
72. 5..004.9
18. 18.13:212
81,90en.21
20.4 M0
68. 5826.4
g-Mixup
20.70=3.30
47.73:4405
67.77a7.4
68.8223344
63.91=2000
GooMia.
70.60=4.01
17. 17.201=5.7
81.90 81,90e7.56
69 10.80=s.s ss
68. 31
GRATIN
71.70e4.44
49.20:3.00 0o
88.83
71.33:5.04
68. 08.6124442
training set to finetune the post-readout function. This filter- ing setup aligns perfectly with the assumptions of Theorem 3.3, as we first train the post-readout function without any augmentation, then evaluate each augmentation' 's influence, and only afterward retrain the post-readout layer using the selected augmented graphs. Our experiments in Section 4 demonstrate that this training paradigm improves general- ization across various datasets and GNN architectures.
delg Ben.m) den.m
I(("T)
[Diest|
A negative value of I(Â))) indicates that adding the aug- mented data to the training set would increase the prediction loss on the test set, negatively affecting the GNN's s gener- alization. In contrast, good augmented graph is one with a positive I( indicating improved generalization. In Figure 2, we present the density of the average influence scores of each augmented data on the test set.
4. Experimental Results
In this section, we present our results and analysis. Our experimental setup is described in Appendix J.
3.5. Fidder-Guided GMM Augmentation
On the Generalization of GNNs. In Tables and 2, we compare the test accuracy of our data augmentation strat- egy against baseline methods. Additional results for the same experiment on larger datasets can be found in Ap- pendix K. We trained all baseline models using the same train/validation/ter splits, GNN architectures, and hyper- parameters to ensure fair comparison. It is worth noting that the baselines exhibit high standard deviations, which is a common characteristic in graph classification tasks. Un- like node classification, graph classification is known to have a larger variance in performance metrics (Errica et al., 2020; Duval & Malliaros, 2022). Overall, our proposed ap- proach consistently achieves the best or highly competitive performance for most of the datasets.
Using influence scores, we can further improve the gen- eralization of the GNN by filtering candidate augmented representations. The process consists of three key stages. (i) Primary GNN training: The GNN model is first trained on the original training set without incorporating any aug- mented graphs. (ii) Augmentation and filtering: A pool of candidate augmented graph representations is generated using a data augmentation strategy based on GMMs. Be- cause computing the gradient Vol(Glest 0) requires access to ground -ruth labels, we evaluate the influence of each candidate augmented graph using the set of validation graph rather than on the unseen test set. This yields a ranking of augmented graphs by their estimated impact on validation performance. During this step, we compute both the gradi- ent and the Hessian only with respect to the post-readout parameters. (iii) Filtering: Finally, we combine a subset of the highest-ranked augmented graphs with the original
Additionally, we observed that the results of the baseline methods vary depending on the GNN backbone, motivating further investigation using influence functions. As demon-
7

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
MUTAG
PROTEINS
DD
40
GCN GIN
15
GCN
250-
GIN
Density
30-
Density
GIN
Density
200-
GCN
10-
100
5
50
0
0-
0
-0.02 0.00 0.02 0.04 0.06 Influence Scores
0.08
0.0
0.1
0.2
0.3
-0.1
0.0
0.1
0.2
0.3
Influence Scores
Influence Scores
Figure 2. The density of the average influence scores of each augmented data on the test set.
Table 3. Robustness against structure corruption: We present the Classification accuracy (# std). The higher the accuracy (in %) the better the model. We highlighted the best data augmentation strategy bold. For this experiment, we use the GCN backbone.
Noise Budget
10%
20%
Dataset
IMDB-BIN
IMDB-MUL
PROTEINS
DD
IMDB-BIN
IMDB-MUL
PROTEINS
DD
DropNode
66.4015.51
44.4612.13
19.1884.
65.7913.23
64.8015.01
43.0612.86
67. 17.73316.3
64.35.14.56
DropEdge
66.7015.10
43.80.3.11
69.3615.90
78.421177
63.2014.30
41.8013.11
68. 18.1005
67.0612.55
SubMix
69.30+3.76
46.73+2.6
69.80+4.73
68.0447.64
63.70+5.64
43.73+3.60
69.09+4.58
19.18+6.29
GcoMix
72.20+5.19
49.20=4.31
70.25
68.00+3.64
70.90+3.85
48.86#5.18
68.3646.01
67.31
G-Mixup
68.30+5.13
45.53+4.12
61.71 +5.81
51, 8.
63. 23.20
44.00=4.63
16.63+5.05
43.71
NoisyGNN
70.50+4.71
40.66+33.12
69.45+-4322
64.18 18
63.50+5.43
38.66+4.12
69.99+3.78
66.24
GRATIN
72.80.2.99
49.36114553
70.61
68.68 13.77
13.10
49. 19.53113.4
70.32
69.0113.00
strated in Theorem 3.3, the gradient, and more generally, the model architecture, significantly influence how augmented data impacts the model's performance on the test set.
scores on GCN, thereby enhancing its performance. This behavior is consistent with the baselines, as most graph data augmentation strategies tend to enhance test accuracy more significantly for GCN than for GIN when applied to DD. This is an interesting phenomenon and worthy of deeper analysis. For the DD dataset, when using the GIN model at inference, we observe Softmax saturation, where the pre- dicted class probabilities approach extreme values (close to 1), c.f. Appendix L. We suspect this saturation to hap- pen due to the large average number of nodes in DD and the fact that GIN does not normalize the node representations, which may lead to a graph representation with large norms. This saturation results in the model making predictions with very high confidence. Consequently, the gradient of the loss function with respect to the input graphs eventually converges to O. In such cases, the influence scores become negligible, explained in Section 3.4.
Robustness to Structure Corruption. Besides general- ization, we assess the robustness of our data augmentation strategy, following the methodology outlined by (Zeng et al., 2024). Specifically, we test the robustness of data augmen- tation strategies against graph structure corruption by ran- domly removing or adding 10% 20% of the edges in the training set. By corrupting only the training graphs, we introduce a distributional shift between the training and test- ing datasets. This approach allows us to evaluate GRATIN's ability to generalize well and predict the labels of test graphs, which can be considered OOD examples. The results of these experiments are presented Table for the IMDB- BIN, IMDB-MUL, PROTEINS, and DD datasets. As noted, data augmentation strategy exhibits the best test accu- racy in all cases and improves model robustness against structure corruption.
Fisher Based Filtering. Figure 3 illustrates the impact of removing augmented representations on test accuracy in the Fisher-Guided GMM Augmentation experiment. The results presented are from a single training run. At the beginning, removing augmented graphs with low or negative influence scores improves generalization. The highest test accuracy reached when significant portion of low-quality augmenta- tions has been removed while retaining high-influence ones. This indicates that a well-selected augmentation subset en- hances model performance. As more augmentations are removed, the overall diversity of the training set decreases.
Influence Functions. In Figure 2, we show the density distribution of the average influence of augmented data sam- pled using GRATIN. These findings are consistent with the empirical results presented in Tables 1 and 2. For the MUTAG and PROTEINS datasets, GRATIN's data augmentation has positive impact on both GCN and GIN models. In contrast, for the DD dataset, GRATIN shows no effect on GIN, while it generates many augmented samples with positive values of the influence
observe that
8

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
Since data augmentation generally helps the model general- ize better, excessive removal reduces its effectiveness. At 100% removal, augmentation is entirely disabled, meaning the model is trained only on the original dataset, i.e., the reference case, leading to a significant drop in accuracy.
also maintains robustness against structural perturbations. An additional advantage of our approach is its efficiency in terms of time complexity, Unlike baselines that gener- ate augmented data for each individual or pair of training graphs, GRATIN fits the GMM to the entire training dataset at once, allowing for fast graph data augmentation without incurring significant additional backpropagation time.
PROTEINS Dataset
0.74
Acknowledgment
Test Accuracy
0.72
Y.A. and M.V. are supported by the French National Re- search Agency (ANR) via the AML-HELAS (ANR-19- CHIA-0020) project. Y.A. and J.L. are supported by the French National Research Agency (ANR) via the "GraspGNNs" JCJC grant (ANR-24-CH23-3888) F.M. ac- knowledges the support of the Innov4-ePiK project man- aged by the French National Research Agency under the 4th PIA, integrated into France2030 KANR-33-RHUS-OA This work was granted access to the HPC resources of IDRIS under the allocation "2024-AD0106134 made by GENCI.
0.70
0.68
0.66
0.64
1.0 Percentage of removed augmented representations
0.0
0.2
0.4
0.6
0.8
Figure 3. Effect of filtering augmented representations test curacy.
Configuration Models. As part of ablation study, we propose simple yet effective graph augmentation strat- egy inspired by configuration models (Newman, 2013). As shown in Theorem 3.1, the objective is to con- trol the term BG-DDÜ-A, [ilbg hhll], which can be achieved by regulating the distance between the original and the sampled graph within the input manifold, i.e.,
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
BG-DØrAA || Si| The approach involves generat- ing a sampled version of each training graph by randomly breaking existing edges into half-edges with probability r and then randomly connecting half-edges until all edges are connected. The strength of this method lies in its simplicity and in preserving the degree distribution. If the distance norm is the L, distance between adjacency matrices, | is an upper bound of E G-D || G where ||] is the average number of edges. The results of this experiment are available in Appendix D.
References
Abbahaddou, Y., Ennadir, S., Lutzeyer, J. F., Vazirgiannis, M., and Boström, H. Bounding the expected robustness of graph neural networks subject node feature attacks. In The Twelfth International Conference on Learning Representations, 2024.
Aboussalah, A. M. and Ed-dib, A. Are gnns doomed by the topology of their input graph? arXiv preprint arXiv:2502.1 17739, 2025.
Aboussalah, A. M., Chi, C., and Lee, C. Quantum comput- ing reduces systemic risk in financial networks. Nature Sci Rep, 13(3990), 2023a.
5. Conclusion
We introduced GRATIN, a novel approach for graph data augmentation that enhances both the generalization and ro- bustness of GNNs. Our method uses Gaussian Mixture Models (GMMs) applied at the output level of the Read- function, an approach motivated by theoretical findings. Using the universal approximation property of GMMs, we can sample new graph representations to effectively control the upper bound of the Rademacher complexity, ensuring improved generalization of GNNs. Through extensive ex- periments on widely used datasets, we demonstrated that our approach not only exhibits strong generalization ability but
Aboussalah, A. M., Kwon, M., Patel, R. G., Chi, C., and Lee, C.-G. Recursive time series data augmentation. In The Eleventh International Conference on Learning Rep- resentations (ICLR), 2023b.
Barshan, E.,, Brunet, M.-E., and Dziugaite, G. K. Relatif: Identifying explanatory training samples via relative in- fluence. In International Conference on Artificial Intelli- gence and Statistics, pp. 1899-1909. PMLR, 2020.

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
Bishop, C. M. and Nasrabadi, N. M. Pattern recognition and machine learning, volume 4. Springer, 2006.
Errica, F., Podda, M., Bacciu, D., and Micheli, A. A fair comparison of graph neural networks for graph classifica- tion. In International Conference on Learning Represen- tations, 2020.
Buffelli, D., Liò, P., and Vandin, F. Sizeshiftreg: a regu- larization method for improving size- seneeaeeeaaliztiinn in graph neural networks. Advances in Neural Information Processing Systems, 35:31871-31885, 2022.
Esser, P., Chennuru Vankadara, L., and Ghoshdastidar, D. Learning theory can (sometimes) explain generalisation in graph neural networks. Advances in Neural Informa- 2021.
Castro-Correa, J. A., Giraldo, J. H., Badiey, M.. and Malliaros, F. D. Gegenbauer graph neural networks for time-varying signal reconstruction. IEEE Transactions on Neural Networks and Learning Systems, 35(9):11734 11745, 2024.
tion Processing Systems,
Fey, M. and Lenssen, J. E. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Repre- sentation Learning on Graphs and Manifolds, 2019.
Garg, V., Jegelka, S., and Jaakkola, T. Generalization and representational limits of graph neural networks. In In- ternational Conference on Machine Learning, pp. 3419- 3430. PMLR, 2020.
Chi, C., Aboussalah, A. M., Khalil, B., Wang. J., and Sherkat-Masoumi, Z. A deep reinforcement learning framework for column generation. In Proceedings of the 36th International Conference on Neural Information Processing Systems (NeurIPS), 2022.
Gaudelet, T., Day, B., Jamasb, A. R., Soman, J., Regep, C., Liu, G., Hayter, J. B., Vickers, R., Roberts, C., Tang, J., et al. Utilizing graph machine learning within drug discovery and development. Briefings in bioinformatics, 22(6))bbab1 159, 2021.
Corso, G., Stärk, H., Jing, B., Barzilay, R., and Jaakkola, S. Diffdock: Diffusion steps, twists, and turns for molecular docking. In The Eleventh International Con- ference on Learning Representations, 2023.
Goodfellow, I., Bengio, Y., and Courville, A. Deep learning. MIT press, 2016
Dabouii, A., Soleymani, S., Taherkhani, F., and Nasrabadi, N. M. Supermix: Supervising the mixing data augmenta- tion. In Proceedings of the IEEE/CVF conference on com- puter vision and pattern recognition, pp. 13794-13803, 2021.
Guo, K., Wen, H.. Jin, W., Guo, Y., Tang. J., and Chang, Y. Investigating out-of- oif-iistribution generalization of gnns: An architecture perspective. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 932 932-943,, 2024.
Dempster, A. P., Laird, N. M., and Rubin, D. B. Maximum likelihood from incomplete data via the em algorithm. Journal of the royal statistical sociely: series B (method- ological), 39(1):122, 1977.
Han, X., Jiang, Z., Liu, N., and Hu, X. G-mixup: Graph data augmentation for graph classification. In Interna- tional Conference on Machine Learning, pp. 8230-8244 PMLR, 2022.
Du, S. S., Hou, K., Salakhutdinov, R. R., Poczos, B., Wang, R., and Xu, K. Graph neural tangent kernel: Fusing graph neural networks with graph kernels. Advances in neural information processing systems, 32, 2019.
Härdle, W., Werwatz, A., Müller, M., Sperlich, S., Härdle, W., Werwatz, A., Müller, M., and Sperlich, S. Nonpara- metric density estimation. Nonparametric and semipara- metric models, pp. 39-83, 2004.
Duval, A. and Malliaros, F. Higher-order clustering and pooling for graph neural networks. In Proceedings of the 31st ACM International Conference on Information & Knowledge Management, pp. 425-435, 2022.
Hong, M., Choi, J., and Kim, G. Stylemix: Separating content and style for enhanced data augmentation. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 14862-148700 2021,
Duval, A., Schmidt, V., Hernández-Garefa, A., Miret, S., Malliaros, F. D., Bengio, Y., and Rolnick, D. FAENet: Frame averaging equivariant GNN for materials modeling. In ICML, 2023.
Huang, Z., Yang, Q., Zhou, D., and Yan, Y. Enhancing size generalization in graph neural networks through disentan- gled representation learning. Forty-first International Conference on Machine Learning, 2024.
Ennadir, S., Abbahaddou, Y., Luizeyer, J. F., Vazirgiannis, M., and Boström, H. A simple and yet fairly effective defense for graph neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 21063-2111 22071, 2024.
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent ker- nel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
10

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
Jagtap, S., Çelikkanar, A., Pirayre, A., Bidard, F., Duval, L., and Malliaros, F. D. Branemf: integration of biological networks for functional analysis of proteins. Bioinform., 38(24):5383-5388 2022.
Morris, C., Kriege, N. M., Bause, F., Kersting, K., Mutzel, P., and Neumann, M. Tudataset: A collection of bench- mark datasets for learning with graphs. In ICML 2020 Workshop on Graph Representation Learning and Beyond (GRL+ 2020), 2020.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization, 2014.
Naim, I. and Gildea, D. Convergence of the em algorithm for gaussian mixtures with unbalanced mixing coefficients. In ICML, 2012.
Kipf, T. N. and Welling, M. Semi-supervised classifica- tion with graph convolutional networks. In International Conference on Learning Representations, 2017.
Nelsen, R. An introduction copulas. Springer. 2006.
Koh, P. W. and Liang, P. Understanding black-box predic- tions via influence functions. In International conference on machine learning, pp. 1885-1894, PMLR, 2017.
Newman, M. Networks: An introduction, 2013.
Newman, M. E., Watts, D. J., and Strogatz, S. H. Random graph models of social networks. Proceedings of the national academy of sciences, (9(uupp 25722, 2002.
Kong, S., Shen, Y., and Huang. L. Resolving training bi- ases via influence-basee data relabeling. In International Conference on Learning Representations, 2021.
Ng, A. Cs229 lecture notes. CS229 Lecture notes, 1(1):1-3, 2000.
Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2012.
Nielsen, F. Cramér-rao lower bound and information geom- etry. Connected at Infinity Selection of Mathematics by Indians, pp. 18-37, 2013.
Law, J. Robust statistics-the approach based on influence functions, 1986.
Ozertem, U. and Erdogmus, D. Locally defined principal curves and surfaces. The Journal of Machine Learning Research, 12:1249-1286, 2011.
Lee, B.-K., Kim, J., and Ro, M. Masking adversarial dam- Finding adversarial saliency for robust and sparse network. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 15126-15136, June 2022.
age:
Panagopoulos, G., Tziortziotis, N., Vazirgiannis, M., Pang, J., and Malliaros, F. D. Learning graph representations for influence maximization. Social Nenwork Analysis and Mining, 14(1):203, 2024.
Li, H., Wang, X., Zhang, Z., and Zhu, W. Ood-gnn: Out- of-distribution generalized graph neural network. IEEE Transactions on Knowledge and Data Engineering. 35 11:7328-3340, 2022.
Pfaff, T. Fortunato, M., Sanchez-Gonzales A., and Battaglia, P. Learning mesh-based simulation with graph networks. In International Conference on Learning Rep- resentations, 2021.
Liao, R., Urtasun, R., and Zemel, {pacc-bayesian approach to generalization bounds for graph neural net- works. In International Conference Learning Repre- sentations, 2021.
Rebuffi, S.-A., Gowal, S., Calian, D. A., Stimberg, F,, Wiles, O., and Mann, T. A. Data augmentation can improve robustness. Advances in Neural Information Processing Systems, 34:29935-29948, 2021.
Ling, H., Jiang, Z., Liu, M., Ji, S., and Zou, N. Graph mixup with soft alignments. In International Conference on Machine Learning, pp. 21335-21344 PMLR, 2023.
Rong, Y., Huang, W., Xu, T., and Huang, Dropedge: Towards deep graph convolutional networks on node clas- sification. In International Conference on Learning Rep- resentations, 2020.
Ma, X., Chu, X., Wang. Y., Lin, Y., Zhao, J., Ma, L., and Zhu, W. Fused grasev-wasserstein graph mixup for geaph-level classifications. Advances in Neural Informa- tion Processing Systems, 36, 2024.
Shalev-Shwartz, S. and Ben-David, S. Understanding mu- chine learning: From theory to algorithms. Cambridge university press, 2014.
Malliaros, F. D. and Vazirgiannis, M. Clustering and com- munity detection in directed networks: A survey. Physics Reports, (33(4):95-142) 2013. Clustering and Commu- nity Detection in Directed Networks: Survey.
Tang, H. and Liu, Towards understanding generalization of graph neural networks. In International Conference on Machine Learning, pp. 33674-33719, PMLR, 2023.

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
Tzikas, D. G., Likas, A. C., and Galatsanos, N. The variational approximation for bayesian inference. IEEE Signal Processing Magazine, 25(6):131-146, 2008.
Vignae, C., Krawezuk, I., Siraudin, A., Wang, B., Cevher, V., and Frossard, P. Digress: Discrete denoising diffu- sion for graph generation. In The Eleventh International Conference on Learning Representations, 2023.
Wang, Y., Wang, W., Liang, Y., Cai, Y., and Hooi, B. Mixup for node and graph classification. Proceedings of the Web Conference 2021, pp. 3663-3674, 2021.
Xu, K., Hu, W., Leskovec, J., and Jegelka, How powerful are graph neural networks? In International Conference on Learning Representations, 2019.
Yang. C., Wu, Q. Wang. J., and Yan, J. Graph neural networks are inherently good generalizers: Insights by bridging GNNs and MLPs. In The Eleventh International Conference on Learning Representations, 2023.
Yang. M.-S., Lai, C.-Y., and Lin, C.-Y. A robust em clus- tering algorithm for gaussian mixture models. Pattern Recognition, 5535033336-- 2012.
Yin, D., Kannan, R., and Bartlett, P. Rademacher complexity for adversarially robust generalization. In International conference on machine learning, pp. 7085-7094, PMLR, 2019.
Yoo, J., Shim, S., and Kang, U. Model-agnostic augmenta- tion for accurate graph classification. In Proceedings of the ACM Web Conference 2022, pp. 1281-1291 2022.
You, Y., Chen, T., Sui, Y., Chen, T., Wang. Z., and Shen, Y.
Graph contrastive learning with augmentations. Advances in neural information processing systems, 33:5812-588 2020.
Zeng, X., Tu, Liu, Y., Fu, X., and Su, Y. Toward better drug discovery with knowledge graph. Current opinion in structural biology, 72:1 12-, 2022.
Zeng, Z., Qiu, R., Xu, Z., Liu, Z., Yan, Y., Wei, T., Ying, L., He, and Tong, H. Graph mixup approximate gro- mov-wasserstinn geodesics. In Forty-first International Conference on Machine Learning, 2024.
Zhang, Y., Li, M., Wang, S., Dai, S., Luo. L., Zhu, E., Xu, H., Zhu, X., Yao, C., and Zhou, H. Gaussian mixture model clustering with incomplete data. ACM Transac- tions Multimedia Computing, Communications, and Applications (TOMM), 17(1s):1-14, 2021.
Zhu, J., Gibson, B., and Rogers, T. T. Human rademacher complexity. Advances in neural information processing systems, 22, 2009.
12

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
A. Proof of Theorem 3.1
In this section, we provide a detailed proof of Theorem 3.1, aiming to derive a theoretical upper bound for both the generalization gap and the Rademacher complexity. Theorem 3.1 Let be a classification loss function with Liip as Lipschitz constant and ((:-) [0, 1l. Then, with probability least 1 ō over the samples Dusin have,
EgND [e(g,0.)] S (aag)) + 5V EgnD E(G, Bmzz)
2 (log(4/0)
2p BaiiheaDD go D. | =|]
N
Moreover, we have,
R(Emu)) S <R()) max Ltip Bm Ax
gm
Proof. We will decompose EgND [e(G, Baug)] EgDD [E()) into a finite sum of terms as follows,
EgnD ((9, (mm))
-EEND [8(),0.) = in + u2 + u3 + 44 u5
where,
41 = Eg-D
[e(g,0mm) B
EsnAs
[(G. 8mnn)
N
[e( (ang)) ]]
((ōmm (ang))
BEND
I Gm
N
m=1
03 N
Σ
BGg= AA
(( 0mes)
1
Σ
E Bg
BGmmAs Ax
[ [00]
N
((@m-,0 0.)]
EGND BsgNAs
[e(3,
04
7m Ax
U5 = EgnD
E BBWAN
[(
0.)
EgND [((G,0.]]
We upperbound each of the terms in the sum. We get,
+ us = EEgnD
{(G )ang)
}]
Egwa
|1
+ EgwD [Ban, [e( - EgwD [((G,0))
< BBnD
[(G 0mmz)
- RBNN
[e(G, 0mng)
]] RewD [Rewas [(5,0.)]] - EgnD [((
2 suu
B
E gmaAN
[ 0]]
EEND [((G,0)) 0)]
Ac0
2sup de
ES-D
E RGmA
[(θ,,0) 0)]
E(G
< 2ssup AE@
REND RS-A
[(() 0) ((9 0)]
| -sl]
< LLip sup EgnDE
A~
DEG
For the term u4, we apply McDiarmid's incquality; we consider the two sets (Gn, H and (G' bD) are identical except at the k-th element, i.e. for a fixed & E (1...., .) we have 9% + Gn and Yn E (1...., ......,,, n # G = G... This setup allows us to bound the change in the expected loss when single graph is replaced. Since the classification loss satisfy ((-) E [0, 1]. we get,
13

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
N
R A~
10(9m: 0)]
1 N n=1
Σ
B S~ax [e(S 0)]
G~aA (e(9n: 0) - e(g 0]]
N
1=1
==1
E Baad [e(gx: 0) E(gh, 0)1
_
N
S 2/N.
The first equality is obtained by the fact that vn + k, that {(+) E 1].
In I'n and Gk # G'k, the last inequality is obtained by the fact
Thus,
N
1
E G~ ( 0.)]
-BND [RenAs [[G,O.)]] 21)
14 0, P(u4 2 2) =
N
n=1
212
I exp
Σm= 4/N2
Nt
exp
Therefore, for d Ei0, 1]. and for = v (log(1/ i.e., exp (-Nx²)
= 6, we have,
P us 2 1oloog11/0 (116)/ 5.
Therefore,
2 (1/5)
2 (1g(1/8)
P < V
= 1 - PP 2
2 - 8. N
N
Thus, with probability of at least 1 - 5,
2 log(1/8)
2 2og (4/g(4/8)
N
N
Moreover, Rademacher complexity holds for u2
2 log (4/)
U2 REGDD RBNAs [(g 0mme))
Σ B [(@m 0mee) S 22((am))) + +41
N
The above inequality tells us that the true risk EG
BBnas [((5, mmmp)) is bounded by the empirical risk
N/ I [((āT 0mmy) plus term depending on the Rademacher complexity of the augmented hypothesis class and an additional term that decreases with the size the sample N.
Additionally, since dang is the optimal parameter for the loss 1 ΣN L8 B: ([ām..0]. thus,
U3 S
By summing all the inequalities, we conclude that,
EgnD [(C) )mmm) EgaD [e(g, 0.)] < 2R((ang) + /5/
2 llgg(46))
+ EgwD As
9th G
N
14

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
Part 2 of the proof.
N
R(eang) - R(() Ren P.
sup
c aug (en: 0)
sup AEE
n en 0)
AEO
n=1
n=1
N
N
S Eg Een
sup
En Lang (9n: 0)
1
Σ
en
Ene(gno 0)
N
N
==1
sup
En ( (n. -
-) e( 0))
BE
N
n=1
E P.
sup BEE
1
Σ
En (Lang (9n, 0) el 0))1
N
N
sup 8eB
La Paug(In: (9ns 0) - ((9n;00)
N
n=1
N
E EG [e(0 0) e(G 0)
sup SEB
N
n=1
Ltip B Bo Ax ||| - -Gnlll-
max
ne N)
B. Proof of Theorem 3.3
In this section, we present the detailed proof of Theorem 3.3, which allows us to perform an in-depth theoretical analysis of our augmentation strategy through the lens of influence functions.
Theorem 3.3 Given a test graph G. from the test set, let 0 arg ming L be the GNN parameters that minimize the objective function in (1), The impact of upweighting the objective function L to L L + ( 0), where G is an augmented graph candidate of the training graph 9n and 6n.m is a sufficiently small perturbation parameter, on the model performance on the test graph gf is given by
delgres Oen.m)
Vol(ghet
0),
den,m
where 0en... arg ming Cang denotes the parameters that minimize the upweighted objective function Lang and H3 (0)( is the Hessian Matrix of the loss w.r.t the model parameters.
Proof. Let TThh be augmented graph candidate of the training graph Gn and 6n.m is a sufficiently small perturbation parameter. The parameters and 0. the parameters that minimize the empirical risk the train set, i.e.,
i arg min L.
0 arg min Ciug arg min L + en, (( 0).
Therefore, examine its first-order optimality conditions,
0 oc
(2)
Oen.m en 0))
(3)
15

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
Using Taylor Expansion, develop Eq. (3). We have lim +0 O Oen.m 0, thus,
VgL(0) m V 0)] + V2c(0) + Gmm Vzel 8)] (om.m -0)
Therefore,
Ó 0enm - -- V [ mn Vjllon 0)]
V VaL(0) + cn.m Vol( 0)
Dropping the ((n,m) terms, and using the Equation 2, i.e. VzL = 0, we conclude that,
mm
= -
V Vgl(O 0).
mn
Therefore,
d8
θ Den.m - = [3
V Vol(C ().
de
den.m
Cn.m
dl(
Den.m)
dl(ghest Pen.m db dden.o
den,m
db den,m
-Vol(( V 0).
C. Mathematical Expressions of GCN and GIN
In this section, we provide concise definitions of two widely used GNN architectures: Graph Convolutional Networks (GCN) and Graph Isomorphism Networks (GIN). These architectures differ in how they aggregate and combine information from neighboring nodes in graph.
Graph Convolutional Network (GCN) (Kipf& & Welling, 2017). The GCN updates node embeddings by aggregating normalized features from their neighbors. Specifically, for node E V, its feature vector (!) layer is computed as,
h
Σ
w(h
UEN(() deg(()
where h (t-1) is the feature vector of node layer t - W(l) is trainable weight matrix for layer and a(.) is non-linear activation function, such as ReLU.
Graph Isomorphism Network (GIN) (Xu al., 2019). GIN is designed match the expressiveness of the Weisfeiler- Lehman (WI.) graph isomorphism test. It updates node's embedding by aggregating its own feature with those of its neighbors, followed by multi-layer perceptron (MLP). The update rule at each message passing layer is
h MLP(() (1+c()) (t + Σ h
h) mEN(w)
where (()) is a fixed or learnable scalar. GIN allows for more expressive feature transformations compared to methods with fixed aggregation schemes, enabling it to distinguish broader range of graph structures.
16

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
Matrix forms and comparison. The node representations h each layer can expressed matrix form as H [exd, where p the number of nodes in the graph, de is hidden dimension in the t-th layer, and H is the concatenation of all node representations h for E V. For GCN, the update rule can be written as:
H (D- 1/2H (11 w(o),
where À + is the adjacency matrix with self-loops, and is its diagonal degree matrix. In contrast, for GIN, the update rule is given by:
H MLP(()
AH
If the same MLP is used, comprising a learnable linear layer followed by ReLU activation, the only difference between GCN and GIN lies in the graph shift operator: GCN uses the degrec-normalized operator D D-1/2 AD 1/22 while GIN uses (1+c I+ A.
Therefore, when node features are taken as constant and only the graph structure considered, the difference in their Lipschitz behavior can be traced back to the function that maps the adjacency matrix (and degrees) to these respective shift operators. In other words, the Lipschitz constant difference arises from whether the adjacency information is normalized (D (D-1/2 D--22/2)
((1 (0) A).
Lipschitz constants of GCN and GIN. Under constant node features, the Lipschitz constant each architecture can be decomposed into two parts: one accounting for the sensitivity the mapping from the adjacency matrix to the respective graph shiff operator, i.c., degreeenormalizedd 1-1/2 A 1/2 for GCN vs. (1 + e)I A for GIN, and the other capturing all shared learnable transformations.
1/2
Let lGCN and IGIN be respectively the graph shift operators A D-1D-1222 A D 112 and A (1 + ++)I+ + A, and let parrass represent the product of Lipschitz factors arising from the shared functions. Then,
LGCN LGGN pparans: LGIN LGIN
Eparams.
becomes clear that the difference between GCN and GIN Lipschitz constants depends solely on whether the graph adjacency is normalized, since Eparams is common both. It's straightforward that fGIN I 1 since the corresponding graph shift operator is just translation. Let us now derive an upper bound for ÉGIN- Let A Ag E IRP be adjacency matrices of graphs on n nodes. For each adjacency matrix A,, we define the diagonal degree matrix, D; diag(deg)) (+) degp) (1)
where Vj S p, deg)) (t)
= Ek (A=)
We have,
D D -D2
A A2 D2
D
A D DD
A D
+ D2
A D
D
A2 D
D
(A A2) D
D
D
A
D
2 + D
A2
D
D
(i)
We can upperhound each of (i), and (ii). For (i), we have the following upperbound,
1 |D] 2 (A1 A²) D D 1
Azl
||A
Azll:
81,min
where 6kmin = min, deg, For the second (ii), if we consider for example |||| the L,1 norm, then,
17

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
[[ < [D] 1/2
D2
1/2
||
1/2
[ 1/211/1111
2
1/2
D/2/2/11 |||
D2 1/2
min (1,1m.mm min; 82 min
1/2
2M
|D 1/22 -2 1/2
min (00001m,.m min 8 min
1/2
2M
[ -D2l
min(6 (in( min m
5/2
2M
A
A21pl
min 01 min 82 min
5/12
2Mp
AA - A2ll,
min min(d1.mmm min m2.m 5/2
where M is the maximum norm of the adjacency matrix in the graph dataset, and 1p E R' the vector ones.
Putting both inequalities together, we can come up with an upperbound for fGON that depends on the minimum degree in the dataset.
D. Configuration Models
In this section, we present a novel adaptation of configuration models as a graph data augmentation technique for GNN. Configuration models (Newman, 2013) enable the generation of randomized graphs that maintain the original degree distribution. We can, therefore, leverage this strategy to improve the generalization of GNNs. Below, we present the steps involved our approach using Configuration Models for Graph Data Augmentation:
1. Extract edges: For each training graph G., we first extract the complete set of edges En.
2. Stub creation: Using Bernoulli distribution with parameter E [0, 1], randomly select subset of candidate edges and break them to create stubs (half-edges).
3. Stub pairing: We then randomly pair these stubs to form new edges, creating randomized graph structure with the same degree distribution.
Table 4 shows the performance of this approach on the two GNN backbones, GCN and GIN.
Table 4. Classification accuracy (# std) different benchmark graph classification datascts for the data augmentation baselines based on the GIN backbone. The higher the accuracy (in %) the better the model.
Model
IMDB-BIN IMDB-MUL
MUTAG
PROTEINS
DD
Config Models w/ GCN
71.70+3.16
48.40=3.88
74.97+6.77
70.08
69.01
Config Models w/ GIN
71.70+4.24
49.00=3.44
81.43+10.05
68.3445.30
71.61
As noticed, the configuration model-based graph augmentation method performs competitively with the baselines and even outperforms them in certain cases. This underscores the importance of Theorem 3.1. When compared to our approach GRATIN, the latter gives better results across different datasets and GNN backbones. This difference primarily due the configuration model based approach being model-agnostic. whereas GRATIN leverages the model's weights and architecture, explained in Section 3.4 and supported by Theorem 3.3.
E. Ablation Study
To provide additional comparison and motivate the use of GMMs with the EM algorithm within GRATIN, we expanded our evaluation to include additional methods for modeling the distribution of the graph representations. Specifically, the comparison includes:
18

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
GMM w/ Variational Bayesian Inference (VBI); We specifically compared the Expeetation-MMxxiiizati (EM) algorithm, discussed in the main paper, with the Variational Bayesian (VB) estimation technique for parameter estimation of each Gaussian Mixture Model (GMM) (Tzikas et al., 2008) for both the GCN and GIN models. The objective of including this baseline is to explore alternative approaches for fitting GMMs to the graph representations.
Kernel Density Estimation (KDE); KDE a Neighbor-Based Method and non-parametric approach to estimating the probability density (Härdle al., 2004). KDE estimates the probability density function by placing a kernel function (e.g., Gaussian) at each data point. The sum of these kernels approximates the underlying distribution. Sampling can be done using techniques like Satropolis-Hastingee The purpose of using KDE as baseline is to evaluate alternative distributions different from the Gaussian Mixture Model (GMM).
Copula-Based Methods: We model the dependence structure between variables using copulas, while marginal distributions are modeled separately. We sample from marginal distributions and then transform them using the copula (Nelsen, 2006).
Generative Adversarial Network (GAN): GANs powerful generative models that learn approximate the data distribution through adversarial process between two neural networks. To evaluate the performance of deep learning- based generative approaches for modeling graph representations, we included IGAN, GAN architecture specifically designed for tabular data (Yang et al., 2012). We particularly train tGAN on the graph representations and then sample new graph representations from the generator.
Table 5. Ablation study on the density estimation scheme for learned GCN representations in GRATIN.
Model
IMDB- BIN
IMDB-MUL
MUTAG
PROTEINS
DD
GMM EM
71.00=4.40
49.82 4.22
76.05
70.97=5.07
71.90
GMM w/ VBI
71.00=4.21
49.53
76.05 66.47
70.97
71.64
KDE
55.90 10.29
39.53+22.7
66. 66.64416.79
59.56+2262
58.66=3.97
Copula
69.80
47.13#3.45
74.44446.26
65.04
65. 65.70
GAN
70.60
48.80#5.51
75.52:44.98
69.98 4.5.46
66.26=3.72
Table 6. Ablation study on the density estimation scheme for learned GIN representations in GRATIN.
Model
IMDB-BIL BIN
IMDB-MUL
MUTAG
PROTEINS
DD
GMM w/ EM
71.70=4.24
49.20 29.20:22.0
88. 28.83 02
71.33:3.04
68 6
GMM w/ VBI
71.40+2.65
27.20
88.3045.19
70.25
67.82=119
KDE
69. 09.10=3..95
41.46=3.02
77.60+4.68
60,3743.04
67.48=6.18
Copula
70.60+2.61
47.60=2.29
88.30+5..19
70.16+4.55
67.91
GAN
70.50+3.800
48.40=1.71
38.83=25.00
71.33=3.555
67.74=1.82
We compare these approaches for both the GCN and GIN models in Tables 5 and 6, respectively. 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.
F. Training and Augmentation time
We compare the data augmentation times of our approach and the baselines in Table 7. In addition outperforming the baselines on most datasets, our approach offers an advantage in terms of time complexity. The training time of baseline models varies depending on the augmentation strategy used, specifically, whether it involves pairs or individual graphs. Even in cases where graph augmentation has a low computational cost for some baselines, training can still cone-conu 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,
19

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
thereby increasing the overall training time. Following the framework of (Yoo al., 2022), introduces augmented graphs for each training graph at every epoch to achieve optimal results. In contrast, GRATIN 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 and 2.
sample several
Table Mean training and augmentation time in seconds our model comparison to other benchmarks.
Time
Model
IMDB-BIN
MUTAG
DD
Vanilla
DropEdge
0.02
0.01
0.01
DropNode
0.01
0.02
0.01
Aug. Time
SubMix
1.27
0.23
0.45
G-Mixup
0.74
0.11
4.26
GeoMix
2,344.12
73.52
1,005.35
GRATIN
2.87
0.51
3.25
Vanilla
765.96
99.32
428.10
DropEdge
892.14
596.82
3,037.30
DropNode
884.71
803.63
3,325
Train. Time
SubMix
1,711.01
1,487.03
2,751.92
G-Mixup
148.71
28.14
177.55
GeoMix
89.01
101.82
123.41
GRATIN
774.47
101.56
438.39
Table Classification accuracy (+: std) different benchmark datasets for the data augmentation baselines using the GCN backbone. Higher accuracy (in %) is better.
Model
IMDB-BIN
IMDB-MUL.
MUTAG
PROTEINS
DD
DropEdge
71.4013.60
48.4712.36
73.8817.98
67.5614.50
66.0414.35
DropNode 71.80
48.6013.45
72.78
67.83
66.55
SubMix
72.5013.90
47.93
71.72:10.59
59.7412.83
62.2513.29
Table Classification accuracy (+: std) on different henchmark datasets for the data augmentation baselines using the GIN backhone. Higher accuracy (in %) is better.
Model
IMDB-BIN
IMDB-MUL
MUTAG
PROTEINS
DD
DropEdge
71.70=3.25
43.67=6.58
72.36 12.36411..0 40
66.31#s.d4
70. 10.20+2..5
DropNode
71.00+4.19
42.6745.54
75.00+6.07
54.
70.80+3.31
SubMix 71.10+4.15
42.8043.93
84.53+6.00 06
60.1043.78
70.03+2.01
Furthermore, the performance of simple augmentation baselines such as DropEdge, DropNode, and SubMix significantly drops when using only augmentation per training graph for all the epochs, which is the framework adopted in GRATIN, as shown in Tables 8 and 9. Tables 8 and 9, these baselines experience significant drop in accuracy under this constraint, emphasizing their reliance on generating diverse augmentations at each epoch to maintain strong performance, c.f. the results in Tables I and 2.
The only graph augmentation baseline with comparable or better time complexity than GRATIN G-Mixup. However, GRATIN consistently outperforms G-Mixup most cases, as shown by the results in Tables and 2.
20

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
G. Augmentation Strategies
We present several graph data augmentation strategies, each defined by the augmentation function Ax. For simplicity, we denote it (9,)) instead of the more explicit (9 yn).
Edge Perturbation. Randomly adding removing edges. Given graph 9n, (Vn, En; X,), edge perturbation is defined as (G,) = (V & U Eadd \ Evemover Xn), where edges in Endd and Eremove are sampled from a Bernoulli distribution P()) (A) with probability X.
Node Feature Perturbation. Augmenting node features by introducing noise or masking some features. This is given by (9,)) = (Vn, En, X,, + XZ), where Z ~ (0, I) for Gaussian noise addition.
Subgraph Sampling. Extracting subgraphs. A common approach is k-hop neighborhood sampling, where (x,)) (V%, &, X,); with Vn S Vn and E, being edges induced by the k-hop neighbors of randomly selected nodes selected based on prior distribution P(A).
GRATIN. In our approach, the augmented hidden representations h& AA (hg) corresponds to sampled vector from the GMM distribution Pa that was previously fit on the hidden representations He = {he I yn c} of the graphs in the training set with the same class c. Formally, h3 = Ax. g E c)), where de are the parameter of the GMM distribution Pe-
It is important to note that in some augmentation strategies, the augmented graphs 9ITh may not explicitly depend on the specific training graph Gn. Instead, they may be sampled or generated based on other factors, such general graph distribution or global augmentation rules. This flexibility allows the augmentation framework to capture broader range of variations while maintaining consistency with the original training data.
H. Graph Distance Metrics
Let us consider the graph structure space (A, I1-I and the feature space (X. Il-lxx),, where |1-|/ and ||-IIx denote the norms applied to the graph structure and features, respectively. When considering only structural changes with fixed node features, the distance between two graphs G, G is defined as
G -
g = =|AA - |
(4)
where A, A are respectively the adjacency matrix of G. and the norm I|-lle can be for example the Frobenius or spectral norm. If both structural and feature changes are considered, the distance extends
e و = =a - ĀlLA + B - X
(5)
where X, X are the node feature matrices of G, 9 respectively, and a, 3 are hyperparameters controlling the contribution of structural and feature differences.
In most baseline graph augmentation techniques, such as G-Mixup, SubMix, and DropNode, the alignment between nodes in the original graph G and the augmented graph & is known. However, in cases where the node alignment is unknown, we must take into account node permutations. The distance between the two graphs is then defined as
c
PEll min (a|A- 124-PPM PPPPPPPT ll + X -XXXI PX
(6)
where 11 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 (Abbahaddou et al., 2024).
A common choice for measuring the distance between two graphs is the norm applied to their adjacency matrices. One widely
used norm is the Frobenius norm, defined as ||A- Ä |F - v Ei, A,,)², which captures element-wise differences
between the adjacency matrices. Another commonly used norm is the spectral norm, defined as ||A - Ál| ||2 = O A), where mmax denotes the largest singular value of the difference matrix.
21

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
I. Gaussian Mixture Models
GMMs are probabilistic models used for modeling complex data by representing them as a mixture of multiple Gaussian distributions. The probability density function p(x) of data point x in GMM with K Gaussian components is given by:
p(x) Σ T&N(x pk Ex),
where Te is the weight of the k-th Gaussian component, with 2 0 and Ek=l =,, and and N(x p& Ex) is the Gaussian probability density function for the k-th component, defined as,
1 N(x Mg Ex) (2T) (2r)d/2 det (Ex)1/2 exp (- (x- p4) T 'x''xx 1n))
1
where by and Σn are respectively the mean vectors and the covariance vectors of the k--th Gaussian component, and d the dimensionality of x. The parameters of a GMM are typically estimated using the EM algorithm (Dempster et al., 1977), which alternates between estimating the membership probabilities of data points for each Gaussian component (Expectation step) and updating the parameters the Gaussian distributions (Maximization step). GMMs are a powerful tool statistics and machine learning and are used for various purposes, including clustering and density estimation (Ozertem & Erdogmus, 2011 Naim & Gildea, 2012; Zhang et al., 2021).
J. Experimental Setup
In this section, we detail the experimental setup for the conducted experiments. The necessary code to reproduce all our experiments is available on github at:
Datasets. We evaluate our model five widely used datasets from the GNN literature, specifically IMDB-BIN, IMDB- MUL, PROTEINS, MUTAG, and DD, all sourced from the TUD Benchmark (Morris al., 2020). These datasets consist of either molecular or social graphs. Detailed statistics for each dataset are provided in Table 10. We split the dataset into train/test/validation set by 10%%//0%//195 and use 10-fold cross-validation for evaluation following the recent work of Zeng et al. (2024). If a dataset does not contain node features, we follow the standard practice in GNN literature by using one-hot encoding of node degrees as input features.
Table 10. Statisties of the graph classification datasets used in experiments.
Dataset
#Graphs
#Features
Avg. Nodes
Avg. Edges
#Classes
IMDB-BIN -BBB
1,000
19.77
96.53
IMDB-MUL
1,500
13.00
65.94
MUTAG
188
7
17.93
19.79
PROTEINS
1,113
3
39.06
72.82
DD
1,178
82
284.32
715.66
COLLAB
5,000
74.5
4914.4
REDDIT-MSK
4,999
508.5
1189.7
Baselines. We benchmark the performance of our approach against the state-of-the-art graph data augmentation strategies. In particular, we consider the DropNode (You et al., 2020), DropEdge (Rong et al., 2020), SubMix (Yoo et al., 2022), G-Mixup (Han et al., 2022) and GeoMix (Zeng et al., 2024). For the robustness to structure corruption experiment, c.f. Table 3, we also included NoisyGNN (Ennadir et al., 2024), a recent method explicitly designed enhance GNN robustness, as an additional baseline.
Implementation Details. We used the Py PoTch Geometric (PyG) open-source library, licensed under MIT (Fey & Lenssen, 2019). The experiments were conducted on RTX A6000 GPU. For the datasets from the TUD Benchmark, we used a size base split. We utilized two GNN architectures, GIN and GCN, both consisting of two layers with hidden dimension of 32. The GNN was trained on graph classification tasks for 300 epochs with a learning rate of 10 using the Adam optimizer (Kingma & Ba, 2014). To model the graph representations each class, we GMM using the EM algorithm, running for 100 iterations or until the average lower bound gain dropped below 10-3 The number of Gaussians used in the GMM
22

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
provided in Table 11, After generating new graph representations from each GMM, fine-tuned the post-readout function for 100 epochs, maintaining the same learning rate of 10-2.
Computation of Influence Scores. Computing and inverting the Hessian matrix of the empirical risk is computationally expensive, with complexity of (N p² p²). where = |0] is the number of parameters in the GNN. To mitigate the cost of explicitly calculating the Hessian matrix, we employ implicit Hessian-vector products (iHVPs), following the approach outlined Koh & Liang (2017).
Table 11. The optimal number of Gaussian distributions in the GMM for each pair of dataset and GNN backbone.
Model
IMDB-BIN
IMDB-MUL
MUTAG
PROTEINS
DD
GCN
10
10
GIN
50
5
2
2
50
K. Experiments on Large Datasets
We conducted additional experiments larger-scale datasets, including COLLAB and REDDIT-MULTI-SK (Morris et al., 2020). The resufffor and GIN backbones are included in Table 12 and 13, which further confirm the effectiveness of our approach on graphs with larger number of nodes and edges.
Table 12. Classification accuracy (+: std) on large datasets for the data augmentation baselines using the GCN backbone. Higher accuracy (in %) is better. The symbol denotes instances where the method has an excessive augmentation time, i.e., exceeding 2 hours
Vanilla
DropEdge
DropNode
SubMix
Mixup
GeuMix
GRATIN
COLLAB REDDIT-MULTT-SK
79.94 + 1.61 48.88 + 2.31
79.70 1.10
79.62 +: 1.84
81.86 + 1.62
81.76 + 1.58
80.74 4: 1.89
82 28 1.82
48.87 1.99
48.73 + 2.39
48.77 + 2.01
46.23 2.74
49.31 + 1.56
Table 13. Classification accuracy (# std) on large datasets for the data augmentation baselines using the GIN backbone. Higher accuracy (in %) better. The symbol denotes instances where the method has an excessive augmentation time, i.e. exceeding 2 hours
Vanilla
DropEdge
DropNode
SubMix
G-Mixup
GeoMix
GRATIN
COLLAB
77.80 + 1.53 REDDIT-MULTI-SK
78.26 1.46
78.86 + 2.09
80.98 1.24
78.89 2.33
78.20 + 1.31
79.08
51.85 + 4.29
44.52 9.58
50.87 # 3.36
49.93 = 3.63
50.63 4.04
51.53 3.54
L. Softmax Confidence and Entropy Distributions
One of the critical challenges in training GNN is Softmax saturation, where the model produces confident predictions. This high confidence leads to vanishing gradients, making the influence score of the augmented graphs converge to zero. To analyze this phenomenon, we examine the Softmax confidence and entropy distributions for GCN and GIN across the different graph classification datasets. In Figure 5, 6, 7 and maximum Softmax confidence and entropy for models trained on the original DD, IMDB-BIN, IMDB-MUL, MUTAG, and PROTEINS datasets. Each dataset panel contains two histograms, the Confidence Histogram and the Entropy Distribution. The Confidence Histogram illustrates the distribution of the maximum Softmax confidence scores assigned by the model to the predicted class. A higher confidence value indicates that the model is more certain about its classification decision. The Entropy Distribution provides a measure of uncertainty the model's predictions, computed as: H(y) = Σ. Me. log(ye), where, for each class c, ye is the predicted probability corresponding to this class. Lower entropy values reflect high- confidence predictions and higher entropy values indicate more significant uncertainty.
present histograms illustrating the distribution of
Specifically, in the DD dataset, we observe extreme case of Softmax saturation in GIN, where almost all predictions collapse to near-man marr-ma confidence. The entropy histogram further reinforces the Softmax saturation in GIN on DD, where entropy values are heavily skewed towards zero, meaning the model rarely assigns significant probability mass outside the predicted class. Compared to other dataset-model settings, this effect is particularly pronounced in GIN trained on DD. In contrast, GCN exhibits a wider confidence distribution, maintaining a more balanced uncertainty.
23

---

GNN Generalization with Gaussian Mixture Model Based Augmentation
IMDB-BIN Dataset
Confidence Hi
Entropy Distribution
ON
Denstty
05
3.6
67
68
6.
10
0.0
43
22
3.3
6.1
6.5
06
03
Maximum Softmax Confidence
Entropy
Figure 4. Softmax confidence and entropy distributions for the IMDB-BII BIN datasct.
IMDB-MUL Dataset
Confidence Histogram
Entropy Distribution
GOI
Density
Maximum Softmax Confidence
Entropy
Figure Softmax confidence and entropy distributions for the IMDB-MUL dataset.
MUTAG Dataset
Confidence Histogram
Entropy Distribution
15.0-
Denslty
12.5-
38.0
2.5
Entropy
Figure 6. Softmax contidence and entropy distributions for the MUTAG dataset.
PROTEINS Dataset
Histogram
Entropy
Dencity
2.7
0.3
0.4
0.5
0.5
Maximum Softmax Confidence
Entropy
Figure Softmax confidence and entropy distributions the PROTEINS dataset.
DD Dataset
Confidence Histogram
Entropy Distribution
w5-
200 1979
ncu
235
Maximum Softmax Confidence
Figure 8, Softmax confidence and entropy distributions for the DD dataset.
24


Reviews:
[{"review_id": "YBmEWzyI4u", "reviewer": "Reviewer_4GM3", "summary": "This paper introduces GRATIN, a novel graph data augmentation approach leveraging Gaussian Mixture Models (GMMs) to enhance the generalization and robustness of Graph Neural Networks (GNNs). The authors argue that GNNs often face challenges in generalizing to out-of-distribution (OOD) data, especially with limited or imbalanced datasets. The proposed method generates augmented graph representations by modeling the distribution of graph embeddings (via GMMs) in the hidden representation space. GRATIN is supported by a theoretical framework based on Rademacher complexity and influence functions, which quantifies the impact of augmentation on generalization performance. Extensive experiments on benchmark graph classification datasets demonstrate GRATIN's effectiveness, achieving competitive or superior results compared to existing augmentation methods.", "questions": "", "limitations": "", "rating": 3, "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": "Response to Reviewer 4GM3\n======\nWe thank Reviewer 4GM3 very much for their careful review. In what follows, we answer point-by-point.\n\n**[S1 and Q1]  Robustness Exp.** While our primary contribution lies in improving generalization via data augmentation, we also address robustness as a secondary but significant aspect, building on the recent framework introduced by the baseline GeoMix. Since, in GRATIN we act on the hidden space generated by several GNN layers, our embedding space jointly considers the structure and features by design (unlike most baselines that focus only on structural perturbations). This comprehensive approach allows our model to significantly outperform baselines under feature perturbation, where others often fail due to their limited augmentation scope. We employed a standard feature perturbation baseline by injecting Gaussian noise $\\mathcal{N}(0, I) $ into node features with a scaling parameter $\\beta=0.5$. We compared GRATIN against a Vanilla GNN (no augmentation) and the recent GeoMix baseline. The results of this experiment, shown in Table 1, demonstrate that GRATIN achieves superior robustness to feature perturbations across both datasets.\n\nTable 1: Robustness to features perturbation\n|Dataset| Vanilla | GeoMix|GRATIN|\n|-|-|-|-|\n|PROTEINS|64.96± 4.08|64.87± 9.28|**69.26±3.52**|\n|MUTAG|62.77±12.40|65.40±8.28|**68.56± 8.81**|\n\n**[S2] Additional Baseline** We appreciate the reviewer’s suggestion. We ran additional experiments using the augmentation method proposed in Luo et al. (ICLR 2023), referred to as GraphAug. This model is applied to a GCN backbone following the same experimental setup used for GRATIN and the rest baselines used in our paper. The results, presented in Table 2, will be included in the camera-ready version of the manuscript. GRATIN outperforms GraphAug across multiple datasets, demonstrating the effectiveness of our augmentation strategy that simultaneously operates over both structure and feature spaces. \n\nTable 2: GraphAug Results\n| Method|IMDB-BIN|IMDB-MUL|MUTAG|PROTEINS|DD|\n|-|-|-|-|-|-|\n| GraphAug |73.91±4.62|49.66±4.22|73.36±9.30|70.80±3.89|71.02±3.84 |\n\n**[S3-Part 1] Sensitivity to GMM Components** We experimented with $K$ components in the GMM in the range from 2 to 50. Below, we include a hyperparameter sensitivity analysis conducted on the GIN backbone using the IMDB-BIN dataset, where we observed consistent results across different numbers of Gaussian components $K$.\n\nTable 3: Impact of $K$ (IMDB-BIN)\n|K|10|20|30|40|50|\n|-|-|-|-|-|-|\n|GRATIN-GIN|71.12±2.70|71.34±2.28|71.38±2.57|71.64±2.72|71.74±4.24|\n\n**[S3-Part 2] Sensitivity to the number of augmented samples**  We conducted a sensitivity analysis to study the effect of the number of augmented samples per graph $m$ on model performance. As shown in the table below, the performance of GRATIN remains stable across a wide range of augmentation levels. This indicates that our method is robust to the choice of this hyperparameter, with accuracy variations within a narrow range even when increasing the number of augmentations from 1 to 30.\n\nTable 4: Impact of $m$ (GRATIN-GCN) \n|$m$|1|5|10|20|30|\n|-|-|-|-|-|-|\n|DD|71.90±2.81|72.25±3.26|72.02±3.34|71.87±3.34|71.81±3.29|\n|MUTAG|76.05±6.74|75.53±6.76|75.53±6.76|75.37±6.54|75.16±6.42|\n\n**[S4] Related Work**  Thank you for pointing out these relevant works. We appreciate the suggestions and will include both references in the camera-ready version of the manuscript. GMMDA focuses on node classification, proposing a GMM-based augmentation that preserves labels through MDL-guided sampling of synthetic nodes. GCA, on the other hand, targets graph community augmentation, generating unseen graphs with new community structures by introducing new clusters in the latent space. Both works support the general motivation of modeling latent graph representations with GMMs, which aligns with our method. However, our approach differs in that it uses GMM sampling not to generate new nodes or community structures but to augment graph-level representations in the hidden space of a GNN. We view these works as complementary and helpful in motivating GMM-based augmentation in GNNs, and will cite them accordingly. \n\n**[Q2] Dynamic Integration of Hyperparameters** In our current setup, certain parameters, such as the number of Gaussian components, were observed to have a negligible impact, whereas others, like the number of augmentations, directly influence the size of the training set and, consequently, the overall training time. The filtering proportion for augmented representations, which we currently fix based on influence-based heuristics, is indeed a good candidate for a dynamic approach. In principle, we could integrate an attention mechanism or adapt the filtering as an Active Learning problem, allowing the model to automatically learn the most informative augmented samples. We see this as a promising future direction for further improving generalization.", "claims_and_evidence": "I think the main claims in the paper are well supported.", "methods_and_evaluation_criteria": "Overall, the paper conducts reasonable evaluations. \n\nIt can be further enhanced:\n\n1) The paper only studies structure perturbation. Is the model robust to feature perturbation?\n\n2) It is suggested to include learning-based automated graph augmentations, such as:\n[1] Luo, Youzhi, et al. \"Automated Data Augmentations for Graph Classification.\" The Eleventh International Conference on Learning Representations (ICLR). 2023\n\n3) sensitivity analysis of key hyperparameters, such as the number of Gaussian components in the GMM or the number of augmented samples.", "theoretical_claims": "I am not able to check the correctness of the theory, because I am not familiar with the model generalization theory.", "experimental_designs_or_analyses": "See Methods And Evaluation Criteria.", "relation_to_broader_scientific_literature": "1) Advancing Mixup-Based Techniques: GRATIN contributes to the literature by leveraging Gaussian Mixture Models (GMMs) to generate augmented graphs in the hidden representation space rather than directly modifying graph structures. This makes the augmentation highly efficient.\n2) Theoretical Focus: While most prior augmentation methods lack theoretical guarantees, GRATIN provides a rigorous analysis of generalization improvements through augmentation using influence functions and regret bounds.", "essential_references_not_discussed": "I suggest discussing GMM-based augmentation for graph representation learning. \n\nSuch as:\n\n[1] Li, Yanjin, Linchuan Xu, and Kenji Yamanishi. \"GMMDA: Gaussian Mixture Modeling of Graph in Latent Space for Graph Data Augmentation.\" IEEE International Conference on Data Mining (ICDM), 2023.\n[2] Fukushima, Shintaro, and Kenji Yamanishi. \"Graph Community Augmentation with GMM-based Modeling in Latent Space.\" IEEE International Conference on Data Mining (ICDM), 2024.", "other_strengths_and_weaknesses": "Strengths:\n\n1) The use of Gaussian Mixture Models (GMMs) for graph data augmentation at the hidden representation level is both innovative and computationally efficient. Unlike traditional augmentation methods that operate directly on graphs (e.g., DropNode, DropEdge, Mixup variants), GRATIN focuses on the latent space, which avoids costly node alignments and allows architecture-specific augmentations.\n\n2) The paper provides a strong theoretical foundation for the proposed method. The use of Rademacher complexity to analyze generalization bounds and influence functions to measure the impact of augmented data on test performance demonstrates a deep understanding of the augmentation problem.\n\n3) GRATIN is evaluated on a diverse set of graph classification datasets (e.g., IMDB-BIN, MUTAG, PROTEINS, DD) using two prominent GNN architectures (GCN and GIN). The results show that GRATIN achieves strong generalization performance and robustness to structural perturbations compared to baseline methods like SubMix, G-Mixup, and GeoMix.", "other_comments_or_suggestions": "N.A.", "questions_for_authors": "1) Is the model robust to feature perturbation?\n2) Could the hyper-parameters introduced in the method be integrated into the augmentation process dynamically during training? Such as the augmented representations filtering proportion.", "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": "", "extra_scores": {}}, {"review_id": "S9V1THa1z7", "reviewer": "Reviewer_eSPE", "summary": "The authors propose a novel graph data augmentation method, GRAFIN, which leverages Gaussian Mixture Models (GMMs) to learn the distribution of hidden representations generated by a trained Graph Neural Network (GNN). The method then augments the training data based on this learned distribution. Furthermore, the authors provide theoretical results that support the correctness of the proposed approach.", "questions": "", "limitations": "", "rating": 4, "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": "We thank Reviewer eSPE for their review. In what follows, we address the raised questions point-by-point.\n\n**[R1] Penultimate Inequality in the Proof of Thm 3.1**\nThank you, indeed, the inequality can be replaced with an equality. We will update the proof in the camera-ready (CR) version.\n\n**[R2, R3, R4, W1] Typos** We are grateful to the reviewer for spotting the typos. We'll correct them in the CR version.\n\n**[R5] Division by Zero** In our framework, $\\delta_{\\mathcal{D}}$ is a discrete distribution, where each $h\\sim\\delta_{\\mathcal{D}}$ has equal probability $1/N$. Consequently, we disregard sampling augmented representations $\\tilde{h}$ such that $Q_\\lambda(\\tilde{h}) = 1/N$, which is highly unlikely to occur in practice. We'll add a technical assumption in Prop. 3.2 to ensure that our choice of $Q_\\lambda$ excludes augmented representations $\\tilde{h}$ for which $Q_\\lambda(\\tilde{h})=1/N$. This assumption formally avoids the risk of division by zero.\n\n**[W2] Large datasets** The requested results are included in Table 1 below and further confirm the effectiveness of GRATIN on graphs with a larger number of nodes and edges.\n\nTable 1: GRATIN for larger datasets. The symbol ‘--’ indicates augmentation time exceeding 2 hours.\n\n||Model|Vanilla|DropEdge|DropNode|SubMix|G-Mixup|GeoMix|GRATIN|\n|-|-|-|-|-|-|-|-|-|\n|COLLAB|GCN|79.94±1.61|79.70±1.10|79.62±1.84|81.86±1.62|81.76±1.58|80.74±1.89|**82.28±1.82**|\n||GIN|77.80±1.53|78.26±1.46|78.86±2.09|**80.98±1.24**|78.89±2.33|78.20±1.31|80.07±1.35|\n|REDDIT-M5K|GCN|48.88±2.31|48.87±1.99|48.73±2.39|48.77±2.01|46.23±2.74|--|**49.31±1.56**|\n||GIN|51.85±4.29|44.52±9.58|50.87±3.36|49.93±3.63|50.63±4.04|--|**52.01±3.54**|\n\n**[W3] Formula Numbering** We chose to number only equations that are explicitly referenced in the text. However, we’ll number all equations in the revised manuscript to improve readability.\n\n**[Other Comments]** We appreciate the detailed suggestions. We'll align figure captions more closely with visual steps and rephrase ambiguous terms to improve clarity.\n\n**[Q1 and W4] Law of Large Numbers (LLN)** We invoke the LLN primarily for theoretical considerations, illustrating that our method’s assumptions hold as $m$ (the number of augmented graphs per original sample) grows large. In practice, however, even small $m$ consistently improve accuracy without incurring a substantial computational cost (see additional experiments varying $m$ in our response to Reviewer 4GM3, c.f. S3-Part 2). Our findings confirm that a modest number of augmentations can already be beneficial, indicating that while the LLN provides a formal theoretical foundation, our approach remains effective when $m$ is limited.\n\n**[Q2] Post-readout Fine-tuning** We split the GNN into two parts: (1) message-passing layers that produce graph-level representations and (2) a shallow post-readout function $\\psi$ that maps these representations to final predictions. After training the message-passing layers on the classification task, we fit a GMM to the resulting graph-level representations. This hidden space becomes the manifold where we perform our augmentation. If we were to update the message-passing layers after fitting the GMM, the structure of the learned representation space would shift, making the distribution modeled by the GMM inconsistent and, thus, degrading the quality of the augmented samples. To avoid this, we keep the message-passing layers fixed and fine-tune only $\\psi$, which is computationally efficient, scalable and capable of adapting to the augmented dataset. Retraining the message-passing layers would increase computational time. Importantly, the test time prediction remains a composition of the fixed GNN encoder and the updated post-readout function, meaning that the GNN weights responsible for producing hidden representations still play a central role in the model's predictions.\n\n**[Q3] McDiarmid’s Inequality** The line break may have introduced confusion. We refer to the standard application of McDiarmid’s inequality by considering two datasets that differ at exactly one index. This setup allows us to bound the change in the expected loss. We'll revise the text to make this explanation clearer in the final version.\n\n**[Q4] Node Classification** GRATIN can be extended for tasks like node classification. While other methods might also be adaptable, such extensions are not always straightforward in their original formulations. The extension follows the same framework but shifts focus to node-level distributions. We train a GNN, fit class-wise GMMs on node embeddings, sample new representations for augmentation, and retrain a shallow classifier. Experiments with the GCN on widely used node classification datasets (Table 2) demonstrate effectiveness beyond graph-level tasks.\n\nTable 2:GRATIN for node classification.\n\n||Cora|CiteSeer|PubMed|CS|\n|-|-|-|-|-|\n|Vanilla|**80.71±0.61**|70.36±0.90|79.78±0.28|89.45±1.25|\n|GRATIN|80.47±0.51|**71.08±0.71**|**79.85±0.22**|**90.82±0.76**|", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes. I checked all of them (main text + supplementary materials) and found the following issues/suggestions:\n1. In the proof of Theorem 3.1, the penultimate inequality of this proof can be replaced with an equality due to the vanishing of Rademacher variables under the absolute value.\n2. In the proof of Theorem 3.1, in the final equation, shouldn't it be $G_n^{m}$​ instead of $G_n^{\\lambda}$​?\n3. It looks like there's a compilation issue related to mathopdh on page 15.\n4. End of page 15: summation from lowercase n=1, not n=N.\n5. Line 766: There appears to be a potential issue when the measures return the same value for distinct arguments, which could result in a division by zero.", "experimental_designs_or_analyses": "Yes, all experiments in the main text seem correct. The obtained results in Table 3 are bolded as the best, but there are no statistically significant differences compared to the other methods.", "relation_to_broader_scientific_literature": "GMM-GDA builds on prior work in graph data augmentation (e.g., GAug and G-Mixup) by using Gaussian Mixture Models (GMMs) to generate synthetic graph data, improving GNN generalization. Unlike previous methods, it provides a theoretical framework using Rademacher complexity to bound generalization error. Its efficiency and theoretical grounding make it a significant step forward in enhancing GNNs with data augmentation.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths:\n\n\n1. The authors propose a novel method to augment graph data.\n2. The paper is well-written and easy to follow.\n3. The authors present a solid theoretical framework that supports their claims.\n4. The proofs of the theorems are well-structured.\n\nWeaknesses:\n\n1. There are typos that need to be fixed to maintain the professionalism of the paper.\n2. The proposed approach should be validated on datasets with a larger number of graphs and larger graphs in terms of nodes and edges, such as COLLAB and REDDIT-MULTI-5K.\n3. Equations should be numbered.\n4. In the proposed approach, the authors rely on the law of large numbers, which applies for sufficiently large m (notation from the paper), where m is the number of generated graphs per graph in the training set. I am not convinced that this reasoning holds when the number of generated graphs is much smaller.", "other_comments_or_suggestions": "Comments / Suggestions:\n1. Page 3: Replace “Augemtation” with “Augmentation.”\n2. Figure 1: The steps presented in the figure’s caption should be marked on the figure for clarity.\n3. Page 16: Correct “us to e perform” to “us to perform.”\n4. Page 16 (redundant repetition): “The parameters θ̂ and θ̂” should be revised to avoid redundancy.\n5. Structure space term clarification:\n  - Original: “The findings of Theorem 3.1 hold for all norms defined on the graph input space. Specifically, let us consider the graph structure space [...]”\n  - Suggested: The term “graph structure space” might be unclear for readers; consider providing a brief explanation or rephrasing for clarity.\n6. Misleading notation clarification:\n  - Original: “[…] where G denotes the space of all possible graphs […]”\n  - Suggested: “G should denote the space of all possible graphs whose distribution matches that of the training/validation/test sets.”\n7. Line 936: “nodes in the graph, d is hidden dimension” should likely be corrected to “nodes in the graph, where $d_t$ is the hidden dimension.”\n\n\n  -", "questions_for_authors": "Questions:\n1. Could you elaborate more on the application of the law of large numbers? Is it still applicable if the number of generated graphs (m) is smaller than one per graph in the training set?\n2. Could you explain why, after fitting the GMM to representations produced by a GNN, only the learnable parameters of the post-readout function $\\psi$ are fine-tuned? If the reason is that you do not want to modify the GNN weights responsible for producing hidden representations, I believe it is worth explicitly mentioning.\n3. I do not understand the last four lines at the end of page 12. Could you please clarify what you mean by introducing indexing?\n4. Can we extend your framework to node classification? Is that possible?", "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": "", "extra_scores": {}}, {"review_id": "6IuYxYhRJi", "reviewer": "Reviewer_wgAh", "summary": "This paper introduces GRATIN, a novel graph data augmentation method using Gaussian Mixture Models (GMMs) to enhance the generalization of Graph Neural Networks (GNNs) for graph classification. It provides a theoretical framework analyzing the impact of augmentation on GNN generalization via Rademacher complexity, designs the GRATIN algorithm to generate augmented data in the hidden representation space, and demonstrates superior performance in generalization and robustness compared to existing methods through experiments on multiple benchmark datasets.\n\n#Update after rebuttal", "questions": "", "limitations": "", "rating": 3, "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": "Response to Reviewer wgAh\n==============\n\nWe thank Reviewer wgAh for the feedback. In what follows, we address the raised questions and weaknesses point-by-point.\n\n***[Weakness 1] Section 3 Clarity*** We appreciate the feedback and will do our best to improve the organization and clarity of Section 3 in the camera-ready version.\n\n***[Weakness 2] Clarifying Topology vs. Attributes*** A topology shift refers to changes in the structure of the graph, such as the addition or deletion of edges, modifications in connectivity patterns, or global properties like graph density. On the other hand, attribute shift refers to changes in node features, such as different feature distributions and scales. Since GRATIN operates in the hidden space learned using several message passing neural network layers, it naturally captures both topological and attribute variations, leading to robust performance against diverse types of shifts. We will further clarify this advantage in the camera-ready version.\n\n***[Question 1] OOD Differences: Graph vs. Node Classification*** The primary distinction between OOD in graph and node classification lies in the level at which distribution shifts occur. In graph classification, each sample is an entire graph, and the task is typically approached as a fully supervised problem, with a label for every training graph. Consequently, OOD shifts tend to manifest as substantial changes in global structure or attributes. For example, training graphs may be relatively small or simple, while test graphs might be larger or more complex. Conversely, in node classification, each sample is an individual node within a large graph, and the problem is addressed via semi-supervised learning, where only a subset of nodes have labels and the rest must be inferred from the graph’s topology and feature structure. Here, OOD shifts can arise locally, such as when training nodes appear in sparse neighborhoods while test nodes appear in denser regions or when training features lie within a certain range and test features experience domain drift. These differences require distinct OOD modeling strategies: node-level tasks, which often rely on partial labels and independent features, must manage  semi-supervised constraints, whereas graph-level tasks, typically with full supervision on entire graphs, emphasize broader structural or attribute shifts across distinct graphs. Moreover, unlike most baseline approaches, our method, GRATIN, is flexible and can be extended to node classification tasks. We have presented the details of this extension, including experimental results, in our response to Reviewer eSPE (c.f. [Question 4]), and we will incorporate this discussion into the camera-ready version of the manuscript.\n\n\n***[Question 2] Line 229 vs. MMD*** The term in Line 229 is the expected distance between a representation of an original graph $\\mathbf{h}$ and an augmented representation $\\tilde{\\mathbf{h}}$. By contrast, Maximum Mean Discrepancy (MMD) aligns full distributions by comparing mean embeddings in a Reproducing Kernel Hilbert Space (RKHS). While both aim to reduce the gap between original and augmented data, MMD enforces a global distribution-level alignment, whereas Line 229 focuses on local instance-level consistency. Despite these differing scopes, they share the overarching goal of encouraging similarity between real and augmented samples. Practically, the term in Line 229 is simpler, computationally cheaper, and better suited for theoretically motivating our data augmentation strategy.\n\n\n***[Suggestion 1] Formula Numbering*** We thank the reviewer for the suggestion. In the current version, we chose to number only the equations that are explicitly referenced in the text. However, we agree that numbering all equations would improve clarity and alignment with Algorithm 1. We will revise the manuscript accordingly and include equation numbers in the camera-ready version to enhance readability.", "claims_and_evidence": "This paper introduces GRATIN, a novel graph data augmentation method using Gaussian Mixture Models (GMMs) to enhance the generalization of Graph Neural Networks (GNNs) for graph classification. It provides a theoretical framework analyzing the impact of augmentation on GNN generalization via Rademacher complexity, designs the GRATIN algorithm to generate augmented data in the hidden representation space, and demonstrates superior performance in generalization and robustness compared to existing methods through experiments on multiple benchmark datasets.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are suitable for the problem.", "theoretical_claims": "The theoretical proofs (e.g., Theorem 3.1 and Proposition 3.2) are sound and well-presented, offering new insights into the impact of data augmentation on GNN generalization.", "experimental_designs_or_analyses": "Yes. I have check the experimental designs.", "relation_to_broader_scientific_literature": "GRATIN's contributions are closely related to existing research on graph data augmentation and GNN generalization.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "**Strengths**\n\n1) The method shows strong performance across multiple datasets,\n\n\n**Weaknesses** \n\n1) The organization of Section 3 is somewhat confusing and detracts from readability.  \n2) The relationship between the model and the Graph needs further clarification, particularly in distinguishing between topology and attribute distribution shifts.", "other_comments_or_suggestions": "1) Key formulas should be numbered to better align with the steps in Algorithm 1.", "questions_for_authors": "1) What are the differences between OOD in graph classification and node classification tasks?\n\n2) What are the differences between the method described in Line 229 and MMD?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: There is a general consensus among the reviewers to accept the paper.
Decision: Accept (poster)