Response to Reviewer wgAh

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

[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.

[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.

[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.

[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.

[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.

We thank Reviewer eSPE for their review. In what follows, we address the raised questions point-by-point.

[R1] Penultimate Inequality in the Proof of Thm 3.1 Thank you, indeed, the inequality can be replaced with an equality. We will update the proof in the camera-ready (CR) version.

[R2, R3, R4, W1] Typos We are grateful to the reviewer for spotting the typos. We'll correct them in the CR version.

[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.

[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.

Table 1: GRATIN for larger datasets. The symbol ‘--’ indicates augmentation time exceeding 2 hours.

[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.

[Other Comments] We appreciate the detailed suggestions. We'll align figure captions more closely with visual steps and rephrase ambiguous terms to improve clarity.

[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.

[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.

[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.

[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.

Table 2:GRATIN for node classification.

Response to Reviewer 4GM3

We thank Reviewer 4GM3 very much for their careful review. In what follows, we answer point-by-point.

[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.

Table 1: Robustness to features perturbation

[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.

Table 2: GraphAug Results

[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$.

Table 3: Impact of $K$ (IMDB-BIN)

[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.

Table 4: Impact of $m$ (GRATIN-GCN)

[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.

[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.