Test-Time Graph Neural Dataset Search With Generative Projection

[Re-Weakness(1)] Computation and Explanations of Multiple Module Components in PGNDS

For the three modular components in the proposed PGNDS—dual conditional diffusion, dynamic search, and ensemble inference—we provide a simple one-sentence explanation for each module to help ease understanding:

• Dual conditional diffusion: a generative process that models test-time distribution shifts via reverse-time diffusion.
• Dynamic search: selecting or generating a refined set of test graphs by optimizing over the entire test distribution.
• Ensemble inference: aggregating information from the original test graphs and refined adapted graphs to obtain optimal test-time inference performance.

We have made efforts to clearly structure our paper (e.g., by providing more preliminary background in Section 2.1 and intuitive motivations for each modular design in Section 2.3). We would be happy to clarify any specific concepts that may affect your understanding of our proposed PGNDS.

Even with these three collaborative modules, we have carefully designed the framework to remain efficient at inference time without introducing substantial overhead. As shown in Table 4, PGNDS achieves comparable or even better runtime than prior baseline methods such as (EERM 2.784s vs. PGNDS 2.244s on the OGBG-BBBP dataset). Moreover, our method does not require GNN fine-tuning, and the diffusion steps are truncated ($T \ll T_{\text{tr}}$, refer to Line 128) with early stopping based on the dynamic search criterion in Eq. (12), which significantly reduces unnecessary computation.

---

[Re-Weakness(2)] Feasibility of Well-Trained GNN and Diffusion Models

PGNDS is designed to operate under the common and practical test-time adaptation setting, where reasonably well-trained models are available when deployed in practice. This is consistent with prior TTA methods such as TENT and GTRANS, which also rely on pretrained models for test-time refinement. We argue that this is feasible and reflects a realistic and standard deployment scenario where well-trained models are commonly used. In this context, PGNDS can serve as a test-time plug-in to further enhance generalization without requiring any model retraining.

---

[Re-Questions(1) & (2)] Code and Intuitive Explanation of Dual Conditional Diffusion and Guidance Terms

We will release the code upon the acceptance of this work, considering that it is the first to address the research question of test-time graph neural dataset search. The implementation is not complicated, as it operates entirely at test time without requiring model retraining.

For an intuitive explanation of dual conditional diffusion and the multiple guidance terms: At a high level, PGNDS uses a reverse-time diffusion process with GNN-provided information as control variables to gradually refine test graphs—i.e., following three conditional guidance terms. Intuitively, they work as follows:

• First, graph structure preservation ($r_{\text{struc}}$) ensures that the refined graph remains close to the original input in structural space, which is important as structural information heavily affects graph properties in test-time graph-level tasks.
• Second, task-specific guidance ($r_{\text{gtask}}$) encourages the refined graph to align with the pretrained GNN’s predictions (using pseudo labels), thereby improving task relevance.
• Third, graph diversity ($r_{\text{gdist}}$) prevents the process from collapsing into trivial solutions, maintaining variety in the generated test-time graph candidates.

During each diffusion step, these constraints jointly influence the latent sampling space, guiding it toward a region in the space that balances test-time graph properties with the joint space defined by [training data & the well-trained GNNs].

---

[Re-Questions (3) ] Robustness to Adversarial Perturbations

Thank you for highlighting this interesting and important research direction. To the best of our knowledge, there is currently no prior work specifically studying adversarial robustness for test-time graph adaptation methods, as these methods are inherently unsupervised and operate under the constraint of not modifying deployed models. The field of test-time graph learning is still at an early stage, where the primary focus remains on improving test-time performance under distribution shift.

Although PGNDS is not explicitly designed for adversarial defense, some components naturally contribute to robustness. For example, during the dynamic search phase, PGNDS generates multiple candidate graphs from the dual conditional diffusion and selects those with better task alignment, which serves as a filtering mechanism against suboptimal test graphs. Exploring the robustness of PGNDS under adversarial or noisy settings is indeed a promising direction for future work, and we appreciate the suggestion.

[Re-Weakness(1)] Computation of Dual Conditional Diffusion and Dynamic Search

While PGNDS introduces dual conditional diffusion and dynamic search, we have carefully designed the framework to remain efficient at inference time without introducing substantial overhead. As shown in Table 4, PGNDS achieves comparable or even better runtime than prior baseline methods such as EERM:2.784s vs Ours PGNDS: 2.244s on Ogbg-BBBP dataset. Moreover, our method does not require well-trained GNN fine-tuning, and the diffusion steps are truncated ($T \ll T_{tr}$ refer to Line 128) with early stopping based on the dynamic search criterion in Eq. (12), which significantly reduces unnecessary computation.

---

[Re-Weakness(2)] Well-trained GNNs and Diffusion Model

PGNDS is designed to operate under the common and practical test-time adaptation setting, where reasonably well-trained models are available when deployed in practice. This is consistent with prior TTA methods such as TENT and GTRANS, which also rely on pretrained models for test-time refinement. We argue that this is not an assumption, but rather a realistic and standard deployment scenario where well-trained models are commonly used. In this context, PGNDS can serves as a test-time plug-in to further enhance generalization without requiring any model retraining during test time.

---

[Re-Other Comments Or Suggestions] Limitations and Future Directions

Thank you for the constructive suggestion. One potential limitation of PGNDS is that it currently focuses on graph-level tasks and relies on task-specific guidance derived from model predictions. For instance, extending the current design to node-level settings requires adapting the task-specific guidance term $r_\text{gtask}$ with careful design, especially to ensure compatibility with the structure and graph diversity constraints in Eq. (10).

Looking forward, the proposed test-time graph neural dataset search paradigm opens up several promising directions, such as (1) adapting to more diverse graph types and GNN types (e.g., dynamic graphs and dynamic GNNs) and (2) integrating PGNDS with continual or online learning frameworks under open-world scenarios. We will include a more detailed discussion on these aspects in the final version.

---

[Re-Questions(1)] Under What Circumstances the Generative Projection Process Might Fail

While PGNDS leverages a stochastic generative diffusion process to refine test-time graphs, this randomness is an advantage of our proposed PGNDS, as it provides a broader search space for learning meaningful test-time graph candidates and does not degrade model performance in practice. One potential limitation that might affect the generative projection process arises under severe distribution shifts, where the test-time graph distribution is highly dissimilar from the training distribution. In such cases, the reverse-time projection may produce low-quality or less informative graphs due to insufficient alignment. Nonetheless, our empirical results (Table 1 and Table 2) demonstrate that PGNDS remains robust across a wide range of realistic test-time scenarios.

---

[Re-Questions(2)] Ensemble Inference Ratio

In our current implementation, we adopt a 1:1 ratio between the original test graphs and the adapted graphs during ensemble inference, as defined in Eq. (13). This fixed ratio is both simple and effective, and consistently yields improved performance across datasets, as shown in Table 2 and Table 3. We are open to solutions like learning or tuning adaptive weights; it is not hard to implement but would introduce extra learning parameters and further complicate the method.

---

[Re-Questions(3)] Distribution Alignment and Generalization Guarantee

While PGNDS uses a diffusion model to project the test-time graph distribution toward the training distribution, the goal is not to perfectly change the test graphs to the training graphs, but rather to map test graphs into a joint space defined by [the training distribution & the pretrained GNNs]. This process preserves the intrinsic properties of test graphs while enhancing their compatibility with the well-trained GNN’s learned decision boundary on the training graphs.

In other words, PGNDS performs a model-aware distribution alignment, guided by three constraints in Eq. (10), to improve generalization. This is also supported theoretically in Proposition 2.3 with Eq. (18) - (20), where we show a bounded deviation in the refined distribution, and empirically in Table 2 and Table 3, where the refined graphs consistently yield better predictions.

We thank the reviewers for highlighting the novelty and practicality of our proposed test-time graph neural dataset search (PGNDS), as well as the clear organization and strong empirical results. Detailed responses regarding the more key concept explaination, the dynamic search process and stopping criterion, and the method’s generalizability across diverse tasks and datasets are provided below.

---

[Re-Weakness-(1)] On Method Built on Multiple Concepts

Thank you for the feedback. We understand that the proposed PGNDS introduces multiple concepts—such as graph diffusion, conditional guidance, and dataset-level search—which may be challenging to follow at first. We provide a simple one-sentence explanation for each concept to help ease understanding:

• Graph diffusion: a generative process that models test-time distribution shifts via reverse-time diffusion;
• Conditional guidance: a mechanism to control the reverse-time diffusion for generating refined test graphs using pretrained GNNs' knowledge;
• Dataset-level search: selecting or generating a refined set of test graphs by optimizing over the entire test distribution.

We have made efforts to clearly structure our paper (e.g., provide more preliminary background in Section 2.1 and intuitive motivations for each modular design in Section 2.3). We would be happy to clarify any specific concepts that may affect your understanding of our proposed PGNDS.

---

[Re-Weakness-(2)] On Relation to Graph Neural Architecture Search (GNAS)

Thank you for the question. Our test-time graph dataset search (GNDS) and GNAS have fundamentally different objectives. GNAS aims to search for optimal GNN architectures (e.g., layers, aggregators), typically during the training stage. In contrast, our method PGNDS focuses on test-time dataset-level graph refinement—searching for optimal test graphs to improve prediction under distribution shift, without modifying the model architecture. Therefore, our work serves a completely different purpose from GNAS.

---

[Re-Weakness-(3)] On Dynamic Search and Stopping Criterion

The dynamic search process in PGNDS iteratively samples refined graphs from the conditional diffusion-based candidate set in Eq. (11) and evaluates them using a test-time rectification objective $\epsilon$ in Eq. (12). For each test graph, we select the adapted version that minimizes this objective across the reverse diffusion steps. To ensure search efficiency and stability, we adopt a dynamic stopping strategy: the process halts early if the objective does not improve over a certain number of patience steps (refer to Lines 267–270 in Section 2.4), thus avoiding exhaustive sampling. The effectiveness of this guided selection is also supported by our empirical study—Table 3 (Idx04 vs. Idx05) in Fig.3 and Fig.4 shows that it consistently identifies high-quality test graphs.

---

[Re-Questions For Authors] On Generalization to Other Graph Types and Tasks

Thank you for the thoughtful question. In this work, we focus on molecular and protein graphs with graph-level tasks, covering six benchmark datasets and two types of graph learning tasks (classification and regression, with nine detailed tasks). This demonstrates that PGNDS is effective across a wide range of real-world graph-level scenarios. We would like to emphasize that the core framework of PGNDS is task-agnostic and generalizable. For example, adapting PGNDS to node classification tasks (e.g., on other graph types, i.e., citation or social networks) would involve redefining the task-specific guidance term $r_{\text{gtask}}$ to reflect node-level objectives, and the conditional diffusion and structural constraints remain applicable. We believe the framework’s modular design makes such extensions straightforward.

[Re-Claims and Evidence (1): Difference between “Test-Time Graph Neural Dataset Search” and “Test-Time Graph Adaptation”]

We thank the reviewer for recognizing our contribution in proposing the novel problem of test-time graph neural dataset search (test-time GNDS). In brief, test-time GNDS can be viewed as a distribution-level extension of traditional test-time graph adaptation, where instead of modifying each test graph individually, we learn a parameterized graph distribution at the dataset level. From a broader perspective, both approaches are data-centric test-time graph manipulation methods. We summarize the key differences in the following table, and we will further clarify this conceptual distinction in the revised version.

Aspect	Test-Time Graph Adaptation	Test-Time Graph Neural Dataset Search
Granularity	Per-graph (adapts each test graph individually)	Dataset-level (searches a distribution over graphs)
Adaptation Objective	Modify a single test graph to improve prediction locally	Generate task-aligned graph distributions toward training domain
Method Type	Direct feature and structure manipulation	Generative modeling + sampling via diffusion
Optimization Target	One graph at a time	Entire test-time graph set as a distribution

---

[Re-Claims And Evidence (2)] Whether test-time distribution projection change the properties of refined test graphs?

PGNDS projects the test-time graph distribution toward the training, but it is not to overly change the intrinsic properties of test graphs. This is achieved via the dual conditional diffusion process with three key constraints in Eq. (10): (1) r_struc: structure preservation; (2) r_gdist: graph diversity control; and (3) r_gtask: task-specific preservation. These constraints ensure that the refined test graphs remain closely related to the original inputs while gaining improved compatibility with the training distribution.

Besides, as shown in Proposition 2.3 and Eq. (20), the deviation from the original distribution is bounded by a controllable term $|\xi|$, guided by GNN-specific knowledge. Moreover, empirical results (Table 3, Fig. 3–4) confirm that removing these constraints significantly degrades performance, validating the role of dual conditional diffusion in preserving test graph properties while improving inference.

---

[Re-Methods and Evaluation Criteria] Applicability to Larger Graphs

While our experiments focus on small molecular and protein graphs, PGNDS is not inherently limited to such settings. We evaluate PGNDS on the large-scale QM9 dataset, which contains 133,885 molecular graphs, demonstrating its applicability to large test sets. As a graph-level method, PGNDS can be applied to graphs with more nodes as well, with scalability primarily depending on the efficiency of the diffusion model.

---

[Re-Experimental Designs or Analyses] Ensemble Inference

As confirmed in our ablation study (Table 3, Idx05: only adapted graphs vs. Idx06: ensemble of original and adapted graphs), ensemble inference consistently outperforms using adapted graphs alone. While the adapted test set is aligned with the training distribution, the test-time distribution learning remains inherently complex. Due to the large search space and potential approximation errors in the generation process [refer to Lines 265-272], relying solely on adapted graphs may not capture all informative patterns. Therefore, we propose an ensemble inference scheme to aggregate complementary information from both original and adapted graphs.

---

[Re-Weakness and Suggestions] Recent methods.

Our work is the first to formulate the problem of test-time graph-level dataset search with generative projection. To the best of our knowledge, no existing method directly addresses this setting. While we include representative baselines from both test-time model adaptation (e.g., TENT) and graph adaptation (e.g., GTRANS), there are currently no more recent or directly comparable methods for our proposed task. We will continue to monitor future developments and further explore this new research direction.

---

[Re-Questions For Authors] On Figure A2

Compared to the originals in (a) and (c), the adapted graphs in (b) and (d) exhibit more chemically meaningful substructures (e.g., double bonds).This does not imply that the adapted graph is a “better” version in a ground-truth sense, but rather that it is the most aligned with the training distribution under a well-trained model, and such alignment is beneficial for improving test-time inference performance.