Test-Time Graph Neural Dataset Search With Generative Projection
Authors: Xin Zheng
Wei Huang
Chuan Zhou
Ming Li
Shirui Pan
Venue: ICML 2025 poster (2025)
Keywords: Test-Time Adaption (TTA); Graph Neural Networks (GNNs); Distribution Shift; Generative Models
Primary area: deep_learning->graph_neural_networks
TL;DR: Generative test-time graph adaptation method with distribution learning

Abstract: In this work, we address the test-time adaptation challenge in graph neural networks (GNNs), focusing on overcoming the limitations in flexibility and generalization inherent in existing data-centric approaches. To this end, we propose a novel research problem, test-time graph neural dataset search, which seeks to learn a parameterized test-time graph distribution to enhance the inference performance of unseen test graphs on well-trained GNNs. Specifically, we propose a generative Projection based test-time Graph Neural Dataset Search method, named PGNDS, which maps the unseen test graph distribution back to the known training distribution through a generation process guided by well-trained GNNs. The proposed PGNDS framework consists of three key modules: (1) dual conditional diffusion for GNN-guided generative projection through test-back-to-training distribution mapping; (2) dynamic search from the generative sampling space to select the most expressive test graphs; (3) ensemble inference to aggregate information from original and adapted test graphs. Extensive experiments on real-world graphs demonstrate the superior ability of our proposed PGNDS for improved test-time GNN inference.

Test-Time Graph Neural Dataset Search With Generative Projection
Xin Zheng' Wei Huang?
Chuan Zhou
Ming Li Shirui Pan'
Abstract
ery (Koh et al., 2024; Zhang et al., 2024a; Zheng al., 2022a;b; 2023c;a; Jin et al., 2022; Zheng et al., 2022c; 2023d; Luo et al., 2023; Zheng et al., 2024b; 2025; Yu et al., 2025; Zhang et al., 2024b). Despite the delicate design and thorough training of GNNs on training graphs (Pan et al., 2023; Wang et al., 2024), these models often face significant performance degradation during test-time inference on test graphs the deployment stage (Wu al., 2022; Liu et al., 2023; Chen al., 2023b; Yu et al., 2023; Huang et al., 2025; Zhang et al., 2025). The main cause of this degradation is the distribution shifts between the training and test graphs, which often occur due to substantial changes in node con- texts, edge connections, and graph-level distributions at test time (Wu et al., 2022; Liu et al., 2023; Chen et al., 2023b; Yu et al., 2023; Zhang et al., 2025).
In this work, we address the test-time adaptation challenge in graph neural networks (GNNs), fo- cusing on overcoming the limitations in flexibility and generalization inherent in existing methods. To this end, we propose novel research problem. test-time graph neural dataset search, which learns parameterized test-time graph distribu- tion to enhance the inference performance of un- seen test graphs on well-trained GNNs. Specifi- cally, we propose generative Projection based test-time Graph Neural Dataset Search method, named PGNDS, which maps the unseen test graph distribution back to the known training distribution through a generation process guided by well-trained GNNs. The proposed PGNDS framework consists of three key modules: (1) dual conditional diffusion for GNN-guided gen- erative projection through back-to-training distribution mapping: (2) dynamic search from the generative sampling space to select the most expressive test graphs; (3) ensemble inference to aggregate information from original and adapted test graphs. Extensive experiments on real-world graphs demonstrate the superior ability of our pro- posed PGNDS for test-time GNN inference.
To tackle the issue of test-time performance degradation caused by distribution shifts, test-time adaptation (TTA) has recently gained attention as promising strategy for en- hancing inference performance on test graphs (Chen et al., 2022; Wang et al., 2022; Jin et al., 2023; Zhang et al., 2024e;d). Typically, TTA on graphs aims to dynamically fine-tune well-trained GNN models or modify test graphs to improve GNN model generalization and inference per- formance. Based on the adaptation objective-whette adapting the model or the lata-existing methods can be classified into two categories: (a) test-time model adapta- tion (Chen et al., 2022: Wang et al., 2022: Zhang et al., 2024e): and (b) test-time graph adaptation (Jin et al., 2023). Specifically, given the unseen test graphs, test-time model adaptation mainly works on updating the well-trained GNN models using the self-supervised learning paradigm, where the primary objective is to optimize or fine-tune the pre- trained GNN model parameters. In contrast, test-time graph adaptation takes a data-centric approach, modifying the test graph data while keeping the well-trained GNN model pa- rameters unchanged.
1. Introduction
Graph neural networks (GNNs) have demonstrated impres- sive success in learning from diverse graph-structured data across a wide range of real-world applications, such as rec- ommender systems, traffic forecasting, and drug discoy-
School of Information and Communication Technology, Grif- fith University, Gold Coast, Australia. "RIKEN Center for Ad- vanced Intelligence Project, Tokyo, Japan. "Academy of Mathe- matics and Systems Science, Chinese Academy of Sciences, Bei- jing. China. 'Zhejiang Institute of Optoelectronics, Jinhua, China: Zhejiang Key Laboratory of Intelligent Education Technology and Application, Zhejiang Normal University, Jinhua, China. Corre- spondence to: Shirui Pan <s.pan griffith.edu.au..
However, these two types of methods, as shown in Fig. (a) and (b), face significant challenges in the real world: Chal- lenge 1: Impracticality of fine-tuning deployed GNNs. Once deployed, fine-tuning GNNs for adaptation is imprac- tical due to online application constraints and the high com- putational cost of updating large models. Challenge 2: Inefficiency of data-centric graph adaptation. While test- time graph adaptation avoids modifying model parameters,
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Test- Time Graph Neural Dataset Search
of
Gra
Well- celin-rr ber
Woi-trateed on Gir
trained on C.. GNN Bi
GNN
Piste Adopted Gru
GNN
GNN ON
TTA
TTA
Adiptted 6'
Adopted set of (%
(a) Test- Tist-Time Model Adaptation
(b) Test- (est-Time Graph Adaptation
(c) Test-Time Graph Neural Dataset Search (Ours)
Figure 1. Comparison among test- time adaptation (TTA) techniques on graphs: (a) test-time model adaptation: requires model fine-tuning -enitunn and uses the adapted GNN for inferenee; (b) test-time graph adaptation: requires individually parameterized graph transformation and uses the original well-trained GNN for inference; and (c) our proposed test-time graph neural dataset search: captures parameterized test graph dataset distribution and uses the original well-trained GNN for inference. The red dashed line indicates the final inference process.
aadividualll-ppraeeteiiedd
existing methods like GTRANS (Jin et al., 2023) are limited in flexiiility and generalization, requiring node- and edge- level adjustments and per-graphi fine-tuning making them inefficient for diverse graph distributions.
shifts, the proposed PGNDS framework consists of three essential phases: (1) dual conditional diffusson, which in- corporates well-trained GNN model-specific information, including task-specific graph representations and labels, as dual conditions guide the score updates in the reverse process of the score-based diffusion model. This phase ensures that the refined test graphs preserve their essential data characteristics while becoming more aligned with the training distribution, leading to more effective test-time in- ference performance. (2) dynamic search, which spans the entire reverse process and leverages the learned test graph distribution to identify the most expressive test graphs, lead- ing to a renewed and adapted test graph set. (3) ensemble inference, which aggregates the representative information from both the original and adapted test graphs, enabling the well-trained GNN to more accurate prediction results on the unknown and distribution-shitted test graph set. As result, the proposed PGNDS could achieve flexible and well- generalized test-time adaptation with the automatic generation of renewed test graphs, enhancing the test-time inference performance in a data-centric manner. In sum- mary, the contributions of this work are listed as follows:
In light of these, propose fresh perspective for test-time adaptation on graphs, where the key learning objective for the data-centric solution can be the answer of following:
Question: Rather than modifying individually parameterized test graphs, can we learn a parameterized test-time graph dis- tribution, to project the entire test graph set distribution back to the training set?
The underlying intuition is that, in the ideal scenario with- out distribution shifts, a well-trained model would perform optimally on unseen test graphs. This implies that if the test graph can be effectively mapped back to the training graph in a data-centric transformation process, the generalization error can be minimized the greatest extent.
To answer this question, in this work, we propose a new research problem, test-time graph neural dataset search, which aims to learn the optimal distribution parameters for unknown test graph datasets, enabling them to generalize effectively on well-trained GNN models during test-time in- ference. Specifically, as shown in Fig. 1 (c), test-time graph neural dataset search transforms unknown and distribution- shifted test graphs into new test graph distribution, tailored to specific GNN model architectures, learning behaviors, and task objectives of well-trained GNNs. According to the learned optimal distribution, this approach automatically generates a new test graph dataset, without modifying the well-trained GNN model parameters but enabling better adaptation performance.
New Research Problem. We introduce a new research problem, test-time graph neural dataset search, to learn the optimal distribution of unknown test graph datasets for effective adaptation on well-trained GNN models during test-time inference.
Graph Data-Centric Solution. We develop a generative projection based test-time graph neural dataset search method, named PGNDS, to learn the optimal test graph distribution driven by a diffusion model. By projecting test graphs back to the training distribution, PGNDS enables effectivve test-time adaptation by automatically generating refined test graphs.
To solve this new research problem, we propose a genera- tive Projection based test-time Graph Neural Dataset Search method, named PGNDS, to learn the optimal test graph distribution driven by a score- based diffusion model, by cap- turing test-back-to- back-to-trainne to-training distribution with a GNN-guided generative projection process. Specifically, given a set of test graphs with unknown labels and potential distribution
Comprehensive Experiments. We evaluate the proposed PGNDS on real-world graph datasets, and extensive ex- perimental results demonstrate its superior test-time adap- tation ability for inference on well-trained GNN models.
2

---

Test- Time Graph Neural Dataset Search
Test graph Co
New test graph Gig
Forward stsppa
GNN Bir
pseeddo predicted
Gre
DIFF Win
Dual conditional guidonce
GNN oir
Test graph set
w Ensemble
Dynamic search
Piee
G
G
New test graph set: (GW
Reverse T-step diffusion
Prediction
$1: Test-Time Graph Noural Dataset Search
Inforence
Figure Overall framework our proposed generative projection based test-time graph neural dataset search method (PGNDS).
Related Work. Our research aligns with prior work on test-time adaptation (TTA) on graph (Jin al., 2023; Zheng al., 2023b; Chen al., 2023a; Jang et al., 2022), which focuses on dynamically adjusting pre-trained models or data to enhance adaptation to test samples. More detailed related work of TTA and other distribution shift related problem can be found in Appendix A.
sion tasks as follows:
0r min LONN (Y Yn). where
(I)
Zu Y, GNN),
Here, Our denotes GNN parameters, Zir E RNNxd di is the output graph embeddings, and Y, {fd f.
fe. ... S E R-x11 denotes the output graph labels predicted by the trained GNNo, According to different graph learning tasks, LaNN can be the cross-entropy loss for graph classification, or the mean squared error loss for graph regression, where Y is the ground-truth graph labels. The well-trained GNN model can be denotes as GNNe; with optimal weight pa- rameters 0
2. The Proposed Method
2.1. Problem Definition
Notations. Given training graph set Gu (G ..., GN } P9., containing N graphs, where each graph G (X; At). where X denotes the node feature matrix, A is the adjacency matrix indi- cating whether nodes are connected or not, yir denotes the graph label, and Pe, is the training graph distribu- tion. Similarly, the test graph set can be represented as G.) Pss containing M graphs,
Graph Diffusion Model Training. We focus on the score-based generative graph diffusion model (Jo al., 2022) to capture the distribution of the training graph set. Given the training graph set Su. the diffusion process can be represented as the trajectory of random variables {GH = (X,: A() hejor..) where G E Gur ~ Pa, and GTT follows a prior Gaussian distribution. Then, the diffusion process solves the following stochastic differential equation (SDE) from the graph to the noise:
Gie = [GL G2 where each graph G - -Xi
le: A without labels and Pg. is the test graph distribution. Typically, during test-time in- ference, there exist unknown and distribution shifts between the training and test graph distributions, ie.,, Ps. * P....
Preliminary. This work focuses on addressing the test- time adaptation problem on graphs. During test time, GNN model has been pre-trained using standard supervised learning approach. As part of our proposed method, which incorporates a generative projection technique, we outline the training procedures for both the GNN model and the graph generative model (i.e., diffusion model) as essential preliminary steps.
dG = (G 1) dt + g dw) GP ~ P&g.
(2)
training graph set
where f(- G - G denotes the linear drift coefficient, g(t) R + R denotes the diffusion coefficient, and w is the standard Wiener process. After such a forward process, which adds noise to the original graph distribution to trans- form it into Gaussian distribution, the reverse-time SDE for diffusion can be presented as:
. GNN Model Training. A GNN model is trained on the training graph set Gu for graph-level classification regres-
dG |f (G,,, 96 Vay log (Gt) di+ + gt dw, (3)
where Pt represents the marginal distribution at the time under the forward diffusion process, w is a reverse- time standard Wiener process, and dt is an infinitesi- mal negative time step. To solving the reverse diffusion
3

---

Test-Time Graph Neural Dataset Search
process Vo po (C)) the time-dependent score-based method (Jo al., 2022) trains Se, (:,t) for node features and So, (:, t) for structures until optimal as following:
2.3. Modular Design
Dual Conditional Difffusion. For an arbitrary Cj E 9ie- its state at time step t is denoted as G with the subseript omitted for simplicity. The dual conditional diffusion phase first conducts the forward perturbation, which slightly perturbs the original test graph set to diversify the graph characteristics of both nodes and structures for better back- to to-training projection. Hence, the test-time forward process can be denoted as:
Sup (G 1) 2 Vx, log Pt (G+) Sdz (G t) R VA. log Pr (G+)
(4)
Hence, 25" and x² are well-trained diffusion model param- eters with GNN backbones. It is important to note that all pre-trained parameters from both GNN and diffusion mod- els will remain frozen throughout the entire procedure of our proposed method. Specifically, the goal of the proposed test-time graph neural dataset search can be described as:
dG = f (Gie: 1) di + (1) C% ~ ~Pe (6)
where {G = (Xier HEj...) following the same pro- cess as Eq. (2) but with T < Tu- This constraint helps prevent over-projection of noise, ensuring that the perturbed test graphs remain within the learning space of the well- trained diffusion model.
Definition 2.1 (Test-time Graph Neural Dataset Search). Given the test graph set 9 {G Ghe G. } and the well-trained G model, test-time graph neural dataset search aims to learn the new adapted test graph set Pl. searched from projected distribution P',, with parame- terized mapping function (-)()) as:
After this, we obtain a set of perturbed test graphs (G%- [XL
along with the original clear test graph Ge. for which are fed into the well-trained GNN;; as:
The ~ (F) where F'g, Fre, GNNo;).
F/c = GNN:; (X
yi GNN:; (X w
A (7)
This step provides the predicted label Vic in the reverse-time diffusion and the pseudo ground -oound-truth label 3 to capture the GNN model-specific and graph learning task-driven information from the well-trained GNN.
In Eq. (5), J-le) denotes the search function with as the certain conditional criterion, and Ω represents the parame- ters of the mapping function that requires be optimized.
2.2. Overview of the Proposed PGNDS
Then, taking this (a) task task-specific information and (b) graph data-centric constraint between the (G Gb) as dual condi- tional guidance, the reverse-time diffusion process with the well-trained DIFF (IFF(w;:.:) can be written as
The overall framework of our proposed PGNDS for test- time graph neural dataset search is presented in Fig. 2. Gen- crally, the framework consists of two stages: S1: test-time graph neural dataset search, which mainly covers two sub- steps: (S1-1): dual conditional diffusion and (S1-2): dy- namic search; S2: test-time ensemble inference. Specifi- cally, for (S1-1) dual conditional diffusion, given a set of test graphs G = [G] G2. .... Gh omitted superscript as an illustration example. It first experiences T-step forward process by adding Gaussian perturbation to both node features and edge connections. Then, the perturbed graph Gh is simultaneously fed into the well-trained GNN model GNN for pseudo label y and the well-trained diffusson model ( (+) for reverse diffu- sion process, where the reverse diffusion process is guided by dual conditional information from GNNe; For (S1-2) dynamic search, given the learned test graph distribution P2G., we can obtain a set of potential new test graphs, guided by task-relevant information from GNN;; Then, a dynamic search module is employed to identify the most expressive test graph by spanning the entire diffusion process, leading to renewed and adapted test graph set. For (S2) ensemble inference, with the adapted test graph set, we introduce an ensemble scheme to aggregate the representative informa- tion from both the original and adapted test graphs. More details of the proposed framework are presented as follows.
xx [x 1) G.. (sw: (G 1) - s
di + 91,d
dA [a((At t) 03.5 (so: (ā %
dt da
(8)
where sfeat and sp satue denote the test-time score rectification functions in terms of node features and graph structure, respectively, and they can be calculated as follows:
........Gh .G }, we take an arbitrary
Feat (Tt. Fic, G G(t) siric S/c A Ao) = B3 r'aric ( to AO).
gguask (V Pie) + Igdia ler Gle),
where a, 8, and are the hyper-parameterrs to control the relative importance of three distinct conditional constraints: rgguss (-), Faaue()), and Fydna()). which guide the generation process from the view of: graph learning tasks, graph data- centric diversity, and graph structures, respectively, for rec- tifying the reverse scores on test graphs. More specifically, these three rectification functions can be defined follows:
Fguusk (Fie F1ee) = log (81c
Fite
(10)
Tstruc A Ao = DMSSE A
A
G Cle -Draw
GGe: Gie
Fgdis
4

---

Test-Time Graph Neural Dataset Search
where .)) denotes the mean squared error, and Drow(;)) denotes the fused Wassersennsr distance for defining a distance metric on the input graph space dur- ing test-timg rest-time on the rell-trained GNN model, which excels at determining the optimal coupling between nodes in a fused metric space.
we could obtain new adapted test graph set as %% [G]2 .... This adapted set is aligned with the training distribution while retaining the maximum character- istics of the test distribution. However, given the inherently complex nature of the test-time graph distribution, it is cru- cial to ensure that the generation process can maximally approximate its ground truth. Hence, even with the inte- gration of multiple conditional guidance mechanisms, the search space for identifying the optimal set of adapted test graphs remains extensive. In light of this, we propose an ensemble scheme to aggregate the information from the original test graph Gi and the adapted test graph Gie We adopt the representation ensemble scheme for the graph regression task and the prediction ensemble scheme for the graph classification task.
Specifically, for task specific rectification function Tgssk it uses pseudo ground-truth label The as the condition to guide the diffusion model in generating a test graph that corre- sponds to Yie, capturing the task-driven test graph distribu- tion, ensuring the generated new test graph can sufficiently preserve the label of the original graph during the reverse- time diffusion. For graph structure rectification function Tstruct it calculates the graph structural distance to enforce that the generated graph structure does not deviate signif- icantly from the original one, so that critical connections between nodes are preserved. For graph diversity recti- fication function rgdis, it leverages global information by capturing the interaction between the graph structure and node features, through maximizing the distance between the generated test graph and the original graph, encouraging diversity to prevent the generated graph from collapsing into the original test graph.
Specifically, given Zue Rie GNNe; (Gu) and zig S GNNØ; (Gi,), Zie and z are graph embeddings. For the graph classification task, the final prediction Yets of the well- trained GNN during the test time is the average confidence of both graphs. For the graph regression task, the final output Yreg generated using the fused graph embeddings before being input into the predictor of the well-trained GNN. Hence, we have:
Dynamic Search. After completing the T-step reverse time of the proposed dual conditional diffusion, the proposed PGNDS captures a new condition-guidde test-time graph distribution F's, for adapting the well-trained GNN models. Unlike the standard reverse diffusion, which retains only the final denoised graph, in this work, we build graph dataset search space that spans all reverse steps within the new test graph distribution, we have:
Yels arg max
2 (iie + s16)
(13)
Yreg = foz (n Tie + (1 17) 2/e),
where C E [1....,C} and C' denotes the number of la- bel classes, and foz (.) is the predictor of the well-trained GNNØ;- The proposed ensemble scheme in our PGNDS enable the automatic selection of how much to weigh the original and adapted test graphs, leading better robustness of the test-time inference the well-trained GNN model.
= - 1- 11
Mit E [T. ~ P'g..
(11)
Now the key is to search with a function [-(e)) according to Eq. (5) for the best-adapted test graphs from the Iearned dis- tribution P2g.. This process is guided by a specific criterion for graph data selection, which is designed by incorporating three rectification constraints as outlined below:
2.4. Theoretical Justification
We provide a theoretical analysis demonstrating the feasibil- ity of our proposed dual conditional diffusion in guiding the effective generation of new test graph sets. This approach ensures that the generated test graphs fall into a mixed dis- tribution PVg that combines the training P& and test graph set distributions Pgg,, Po: from GNNe;, and IP (1a ( from
= min aa c' Egtask + B r strue 7
(i.f)
(i.t)
(5.1)
Fgdist
(E(T.
(12)
That means, for each test graph G (1.)) we select the opti- mal adapted graph with the smallest test-time rectification
DIFF (IFFF(WG (#G)
Proposition 2.2. Given the training graph distributionn Ps and the test graph distribution PG.., the complete forward- reverse diffusion process begins with G for T < Tūr steps and ends with potential pseudo back-to-training test graph distribution G ~ PP This process corresponds to a multiple integral of mulli-step transition probabilities, lacking a dlosed-form solution and without guaranteeing distribution aligument.
effort on three-fold constraints at a certain t-step as G/è Additionally, to enhance the efficiency of the search process, we introduce a dynamic stopping condition for the reverse diffusion. Specifically, suppose the overall distance does not decrease after a certain number of patience steps. In that case, the reverse diffusion process is halted, and the best-adapted graph with the current minimum c'.".
Ensemble Inference. After the dynamic search from the learned test-time graph distribution with J(PL lo),
Proof. We primarily use the node feature X ~ Pg.. for
5

---

Test-Time Graph Neural Dataset Search
Table ROC-AUC performance (f) comparison between baseline methods our proposed PGNDS on molecular and protein graphs for the graph classification task, Best results are in bold, and the second bests are underlined.
Molecular
Protein
Methods
Ogbg-BBBP
Ogbg-BACE
Ogbg-ClinToo
Enzymes
ERM (Wu et al., 2022)
0.6935
0.8265
0.6271
0.7161
EERM al., 2022)
0.6857
0.7738
0.6314
0.5274
TENT (Wang et al., 2020)
0.6988
0.8338
0,6254
0.7187
DropEdge (Rong al., 2019)
0.6950
0.8139
0.6256
0.7137
FeatMask (You et al., 2020)
0.6973
0.8240
0.6298
0.7115
GTRANS (Jin et al., 2023)
0.6921
0.8256
0.6312
0.5037
PGNDS (Ours)
0.7014
0.8873
0.6371
0.7245
Table 2. RMSE performance (L) comparison between baseline methods and our proposed PGNDS on molecular graphs for the graph regression task. Best results bold, and the second- are underlined. GTRANS method on QM9-a1pha indicates resulted RMSE exceeding 1000+ signifying unstable upiimization.
QM9
Methods
Ogbg- Free-rr
A
B
C
alpha
ERM (Wu et al. 2022)
2.2022
0.4415
0.3924
4.3220
1.9308
DropEdge (Rong
2019)
2.5879
0.5439
0.3978
9.8058
2,4434
FeatMask (You et al., 2020)
2.3560
0.5183
0.3991
8.3528
2,6432
GTRANS (Jin al., 2023)
2.8192
5.3664
0.5183
1.9289
PGNDS (Ours)
2.1985
0.4283
0.3926
4.2881
1.9252
proof illustration (simplified as Xg ~ g(xn) in the follow- ing), with the structure adhering to the same principles. Then, the forward process can be described as:
For simplicity, we treat X x, for effectively moving the notation x/. back to xT and projecting single Xo. Theoretically, this is a nested Markov chain with multiple integrals of multi-step transition probabilities. The obtained pseudo test distribution 12y differs from the original (i.e., Pa..) or Pv (i.e., i, it is a hybrid product, where Fç, is used the forward direction and Pg,, in the reverse direction. If P.s * PFg, and DIFF & is not perfectly optimized, we can conclude that the final distribution I' (PP IP
Xo ~ g(xo), X1 q(x1 .... XT ~ q(TT) XT-
(14)
Thus, step T, the distribution of x is:
(r(xT) = - - ((x)) II q(xx ... dx (15)
Next, instead of using the reverse process matching g. we employ reverse transformation derived from the well- trained diffusion model DIFF;; with score-based approx- imation, denoting as Py (xt-11 (x)=1/x()?=_r. This reverse sam- pling takes Xy step-by-step back Xg- Thus, the complete forward-reverse combination can be viewed as a unified Markov chain within combined 27 steps:
Given the absence of a closed-form solu- tion and the high degree of freedom of the hybrid pseudo test distribution P. without any guidance, the proposed dual conditional guidance in PGNDS could effectively mitigate distribution misalignment through F by constraining the deviation within a controllable bound &, encapsulated by the well-trained GNN parameter distribution Po,,
Proposition 2.3.
Xg
2
9
XT
XT-1
X/- (16)
Proof. After introducing guidance terms, the reverse diffu- sion Pw (x(_a1=( in Eq. (17) can be modified as:
Here, the final output X may not necessarily equal the original Xg. Then, its distribution can be expressed as:
Pa (x6) 1 f... (..(((o))
[[ g(x: xx
P( o (x/- |x/) Pw (x(x(11xx)
exp (a Tennk
Issrue (x)) + B rgdist (x())
(18)
x 1-T
Pots x
5(x6 - xx) dxn dxi, dxy dodr ... ....
Then, during each step the reverse process, guidance terms apply gradient constraints the generated distribu-
(17)
6

---

Test- Time Graph Neural Dataset Search
tion:
Table 3. Modular indexes for ablation study in our PGNDS.
v Voggpt..e; = V Vlog Po (x/=1xx()
Dual Conditional Diffusion
Dynamic Search
Ensemble
Modules
E Vr. (x.)- (19)
Guidance
Faaakk
Tamu:
Fuik
Inference
Idx01
x
x
Idx02
x
x
Hence, given w+)) represents the distribution discrepancy measurement function (e.g., Wasserstein distance). By incor- porating the guidance terms, the bias of the integral formula can be expressed as reduction in distribution bias as:
Idx03
/
x
Idx04
/
/
/
Idx05
/
/
V
Idx06 (Ouis)
/
(IP', C (().)))) S (P' S.. 77 1116,,
where
for node-level graph OOD generalization, and test-time training method TENT (Wang et al., 2020); And the recent SOTA graph data-centric method: test-time graph trans- formation method GTRANS (Jin et al., 2023). We use the classic GIN model (Xu et al., 2018) as the well-trained GNN due to its widespread use in molecular graph learning.
T
glask + yr + (rrgss)) )|²
(20)
[<| =
V a
1=1
where 7) > 0 is the learning rate and |sl is the controllable bound under our proposed guidance terms.
3. Experiments
3.2. Experimental Results
Test-time GNN Inference Performance. In Table 1, we compare the ROC-AUC performance of several baseline methods with our proposed PGNDS model on both molec- ular and protein graph datasets for the graph classification task. In general, we can observe that our proposed PGNDS achieves the highest ROC-AUC across all datasets, indi- cating the best test-time inference performance for both molecular and protein graphs. Notably, the improvement of PGNDS is significant on the Ogbg-BACK dataset, with a ROC-AUC increase from 0.8338 to 0,8873 over the second- best method. Moreover, TENT provides the second- best re- sults on three datasets except for Tobg-ClinTox. As TENT is a tne-tuning--aat based method for adjusting well-trained model parameters, such results demonstrate that the entropy- based fine-tuning can be beneficial for test-time graph clas- sification tasks. Additionally, EERM and GTRANS out- perform the other baseline methods on the Ogbg-Cling -linnTxx dataset, with EERM securing the second-best result. This can be attributed to its need to adjust the pre-training process to improve GNN generalization ability.
We verify the effectiveness of the proposed PGNDS in terms of the test-time inference performance. Concretely, we aim to answer the following questions: Q1: How does the proposed PGNDS perform on the well-trained GNN for both graph classification and regression tasks when faced with unknown graph distribution shifts at test time? Q2: How does the proposed PGNDS perform in ablation stud- ies focusing on each components? Q3: How sensitive is the proposed PGNDS to variations in hyper-parameters? Q4: How does the proposed PGNDS perform in terms of running time efficiency?
3.1. Experimental Settings
Datasets & Metrics. We perform experiments six real- world graph datasets covering protein and molecular graphs, with four graph classification tasks and five graph regres- sion tasks. We use the area under the ROC curve (ROC- AUC) to evaluate the graph classification task and root mean square error (RMSE) for the graph regression task. Higher ROC-AUC (f) and lower RMSE (4) indicate better graph learning performance. More details of datasets are listed in Appendix B. Note that the original QM9 dataset con- tains nineteen tasks, but we selected four of them (i.e., A, B, C, and a 1 pha) for our experiments. For all training and test graphs, we follow the process procedures and splits in previous works (Liu et al., 2024; Jo et al., 2022).
In Table 2, we compare the RMSE performance of vari- ous baseline methods with our proposed PGNDS model on molecular graph datasets for the graph regression task. Our PGNDS demonstrates superior performance across all datasets, consistently achieving the lowest RMSE values, which indicates its effectiveness in minimizing regression errors. Specifically. PGNDS shows significant improve- ments on the QM9-A and QM9-alpha tasks, achieving RMSE reductions from 2.2022 to 2.1985 and from 4.3220 to 4.2881, respectively. This highlights the robustness of PGNDS in handling diverse molecular graph regression tasks. Among the baseline methods, DropEdge shows com- petitive performance, securing the second-best results on QM9-B, QM9-C, and QM9-a apha tasks, while FeatMask
Baseline Methods & Implementation. We compare the proposed PGNDS with the following baselines that fall in two groups: gruph model-centric methods: empirical risk minimization (ERM) for standard training (Wu et al., 2022), data augmentation technique DropEdge (Rong et al., 2019) and FeatMask (You et al., 2020), Risk Minimization (EERM) (Wu et al., 2022) customized
7

---

Test Graph Neural Dataset Search
Hyperparameter Analysis (d=1)
Ogbg-BBBP
Ogbg-
0.01
euet
Eub
Eu0
8:75
1.00
0.1
4.0
8:70
3.90
455
3,00
6:1
0.90
6:73
6:06
9:44
644
8a5
2.50
E/B
0:00
6.6
0.80
2,00
0,5
18s01
Mesi
Maderterg
14e1
Mao1
Haas
85
6:08
0.90
0:00
6.49
6.43
Abolation Index
Indes
0.01
6.1
0.2
0.5
5
Figure 3. Ablation study results for graph classification (ROC-AUC) on Ogbg-BBBP.
Figure 4. Ablation study results for graph regression (RMSE) on Oghg-FreeSoo FreeSolv.
Figure 5. papeeeparamettr sensitivity analysis for B and y on Ogbg-BBBP with a = 1.
with a ROC-AUC of approximately 0.70. In contrast, mod- els with fewer submodules show a lower ROC-AUC, For the Ogbg-FreeSoly dataset, the RMSE values show simi- lar trend. The model without any guidance has the highest RMSE, indicating poor test-time regression performance. As more components are integrated, the RMSE decreases, with the full model achieving the lowest RMSE.
Table Running time (in seconds) comparison on graph classifi- cation task in 5 epochs with single NVIDIA A100 GPU.
Methods
Ogbg-Bi BBBP
Ogbg-BACE
Ogbg-ClinTox
BERM
2.784
2.287
2,224
GTRANS
0.240
0.224
0.237
PGNDS (Ours)
2.244
2.635
2.269
Hyper-parameter Sensitivity Analysis. In Fig. 5, we an- alyze the hyper-parameter sensitivity for exploring the ef- fects of and y model performance with a 1 for ROC- AUC of Ogbg- BBBP in the dual conditional diffusion pro- cedure for three constraints on Fgunkk Ttrue and Tgdist- The performance is indicated by values across a grid, ranging {0.01, 0.1, 0.2, 0.5, 1, 5). More results for 3 = 1 and y = are presented in Appendix C. The result shows that mod- erate values of 3, and y generally lead to better test-time inference performance. Extreme values, particularly large 3 or very small 7, tend to degrade performance slightly.
performs well on the Ogbg-FreeSoll dataset. We would like highlight that, compared to the graph classification task, test-time graph regression is often more challenging due to the need for precise fitting of specific target graph properties for regression. This difficulty explains why many test-time methods struggle to deliver satisfactory inference performance, particularly in data-centric approaches that revise the test graphs. Optimizing within such large latent space generate specific attributes remains a complex and demanding task.
Running Time Comparison. Table 4 compares the running time of different methods on the graph classification task for first-batch test graphs. Our proposed PGNDS demon- strates competitive efficiency, achieving lower runtime than EERM on Ogbg-BBl -BBP and similar performance on Ogbg- ClinToo. While GTRANS remains the fastest due to its constrained learning space with full-parameter transforma- tion of node features and structure, our method balances efficiency and model complexity by integrating test graph distribution remapping, making it an effective graph-level test-time adaptation method. More visualization results of PGNDS generated test-time graph are listed in Appendix D.
Ablation Study of PGNDS. In Table 3, we list the ablation study settings for evaluating the contributions of different submodules of our PGNDS model. The ablation study systematically removes components from the full model (our PGNDS for Idx06). such as using standard diffusion without guidance (Idx01), the three constraints (Fgusk. Faruee and gddist in the dual conditional diffusion stage (Idx02 to Idx04), dynamic search (Idx05), and ensemble inference (Idx06). Moreover, it is important to note that these mod- ules interact with one another, and thus, the experimental results do not exhibit linear improvement as the compo- nents are added back, highlighting the synergistic effects between the modules. The results are based on ROC-AUC performance for the Ogbg- BBBP classification shown in Fig. 3, and RMSE performance for the Oreg-FreeSolv re- gression shown in Fig. 4, respectively. Performance on Ogbg-BBBP BBBP demonstrates that as more submodules are activated, the performance steadily improves. The best per- formance is achieved by Idx06, the full PGNDS model,
4. Conclusion
In this work, we address the challenge of test-time adap- tation in GNNs by introducing the novel problem of test- time graph neural dataset search. Our proposed method, PGNDS, learns parameterized test-time graph distribution
8

---

Test-Time Graph Neural Dataset Search
to improve inference performance. Leveraging dual condi- tional diffusion, dynamic search, and ensemble inference, PGNDS captures unseen test graph distributions through a generative projection approach. Extensive experiments validate PGNDS superior adaptation capability for graph- level test-time GNN inference. Future research could extend PGNDS to handle dynamic graph neural networks, further broadening its applicability.
android apps. In 2021 IEEE/ACM 43rd International Con- ference on Software Engineering: Software Engineering in Practice (ICSE-SEIP), pp. 101-110. IEEE, 2021.
Huang. Y., Zhang, Z., Zhao, Q. Yuan, X., and Chen, C. Themis: Towards practical intellectual property protec- tion for post-deployment on-device deep learning models. In 34th USENIX Security Symposium (USENIX Security 25), 2025.
Acknowledgment
Jang, M., Chung, S.-Y., and Chung, H. W. Test-time adapta- tion via self-training with nearest neighbor information. In International Conference Learning Representations (ICLR), 2022.
Shirui Pan was supported in part by the Australian Re- search Council (ARC) under grants FT210100097 and 5P240101547 and the CSTRO - National Science Foun- dation (US) AI Research Collaboration Program. Chuan Zhou acknowledged the support from the National Natural Science Foundation of China (No. 62472416). Ming Li acknowledged the support from the National Natural Sci- ence Foundation of China (No. 62172370) and the Jinhua Science and Technology Plan (No. 2023-3-003a).
Jin, M., Li, Y--.,, and Pan, S. Neural temporal walks: Motif-aware representation learning on continuous-time dynamic graphs. In Advances in Neural Information Pro- cessing Systems (NeurIPS), 2022.
Jin, W., Zhao, T., Ding. J., Liu. Y.. Tang. J., and Shah, N. Empowering graph representation learning with test-time graph transformation. In International Conference on Learning Representations' (ICLR), 2023.
Impact Statement
This paper presents work whose goal is advance the field of graph machine learning, specifically in improving the test-time graph inference performance of graph neural net- works (GNNs) in data-centric manner. By proposing a novel test-time graph neural architecture search method, this research seeks to address challenges in deploying GNNs for practical test-time inference. There are many potential societal consequences of our work, none of which feel must be specifically highlighted here.
Jo, J., Lee, S., and Hwang, J. Score- barre-baeddde generative modeling of graphs via the system of stochastic differen- tial equations. In International conference on machine learning, pp. 10362-10383. PMLR, 2022.
Ju, M., Zhao, T., Yu, W., Shah, N., and Ye, Y. GGaah- patcher: mitigating degree bias for graph neural networks via test-time augmentation. Advances in Neural Informa- tion Processing Systems, 36, 2024.
Koh, H. Y., Nguyen, A. T., Pan, S., May, L. T., and Webb, G. I. Physicochemical graph neural network for learn- ing protein-ligand interaction fingerprints from sequence data. Nature Machine Intelligence, pp. 1-15, 2024.
References
Chen, G., Zhang, J., Xiao, X., and Li, Y. Graphtta: Test time adaptation on graph neural networks. arXiv preprint 2rX0v:2208.09126 2022.
Liang. J., Hu, D., and Feng, Do we really need to access the source data? source hypothesis transfer for unsuper- vised domain adaptation. In Internationall Conference Machine Learning (ICML), pp. 6028-6039. PMLR, 2020.
Chen, L., Zhang, Y.. Song, Y., Shan, Y., and Liu, L. Im- proved test-time adaptation for domain generalization. In Proceedings of the IEEE/CVF Conference Com- puter Vision and Pattern Recognition (CVPR), pp. 24172- 24182, 2023a.
Liang, J., He, R., and Tan, T. A comprehensive survey on test-time adaptation under distribution shifts. Interna- tional Journal of Computer Vision, pp. 1-34, 2024.
Chen, Y., Bian, Y., Zhou, K., Xie, B., Han, B., and Cheng, J. Does invariant graph learning via environment augmenta- tion learn invariance? In Advances in Neural Information Processing Systems (NeurIPS), 2023b.
Liu, G., Inae, E., Zhao, T., Xu, J., Luo, T., and Jiang, M. Data-centric learning from unlabeled graphs with diffu- sion model. Advances in neural information processing systems, 36, 2024.
Huang, Y. and Chen, C. Smart app attack: hacking deep learning models in android apps. IEEE Transactions Information Forensics and Security, 17:1827-1840, 2022.
Liu, Y., Ding, K., Liu, H., and Pan, S. Good-d: On unsuper- vised graph
Huang, Y., Hu, H., and Chen, C. Robustness of on-device models: Adversarial attack deep learning models on
detection. In Proceedings of the Sixteenth ACM International Conference on Web Search and Data Mining (WSDM), pp. 339 339-347,,0 2023.

---

Test-Time Graph Neural Dataset Search
Luo, L.,, Li, Y--F, Haffari, G., and Pan, S. Reasoning on graphs: Faithful and interpretable large language model reasoning. In International Conference on Learning Rep- resentations (ICLR), 2023.
Zhang, H., Wu, B., Yuan, X., Pan, S., Tong, H., and Pei, J. Trustworthy graph neural networks: Aspects, methods and trends. Proc. IEEE, 112(2)97-139, 2024a.
Zhang. H., Yuan, X., and Pan, S. Unraveling privacy risks of individual fairness in graph neural networks. In ICDE, pp. 1712-1725. IEEE, 2024b.
Pan, J., Liu, Y., Zheng, Y., and Pan, S. Prem: A simple yet effective approach for node-level graph anomaly detec- tion. In 2023 IEEE International Conference on Data Mining (ICDM), pp. 1253-1258. IEEE, 2023.
Zhang. H., Wu, B., Yang, X.. Yuan, X., Liu, X., and Yi, X. Dynamic graph unlearning: general and efficient post-processing method via gradient transformation. In wwW. pp. 931-944 ACM, 2025.
Rong. Y., Huang, W., Xu, T., and Huang, J. Dropedge: Towards deep graph convolutional networks node clas- sification. In International Conference on Learning Rep- resentations (ICLR), 2019.
Zhang. J., Wang. Y., Yang. X., Wang, S., Feng, Y., Shi, Y., Ren, R., Zhu, E., and Liu, X. Test-time training on graphs with large language models (Ams)). In Proceedings of the 32nd ACM International Conference on Multimedia, pp. 2024c.
Sun, Y., Wang, X., Liu, Z., Miller, J., Efros, A., and Hardt, M. Test-time training with self--uupevision for generaliza- tion under distribution shifts. In International Conference on Machine Learning. pp. 9229-9248. PMLR, 2020.
Zhang. J., Wang, Y., Yang. X., and Zhu, E. A fully test-time training framework for cun--super node classifica- tion on duitt--distiibution graphs. ACM Transactions on Knowledge Discovery from Data, 2024d.
Wang, D., Shelhamer. E., Liu, S., Olshausen, B., and Darrell, T. Tent: Fully test-time adaptation by entropy minimiza- tion. In International Conference on Learning Represen- tations (ICLR), 2020.
Wang, I.., Zheng, Y., Jin, D., Li, F., Qiao, Y., and Pan, S. Contrastive graph similarity networks. ACM Transactions on the Web, 18(2):1-20, 2024.
Zhang. Z., Liu, M., Wang, A., Chen, H., Li, Z., Bu, J., and He, B. Collaborate to adapt: Source-free graph do- main adaptation via bi-directional adaptation. In Proceed- ings of the ACM Web Conference 2024, pp. 664 675, 2024e.
Wang, Y., Li, C., Jin, W., Li, R., Zhao, J., Tang, J.. and Xie, X. Test-time training for graph neural networks. arXiv preprint ar arXiv:2210 10.:2200 2022.
Zheng, X., Liu, Y., Pan, S., Zhang, M., Jin, D., and Yu, Graph neural networks for graphs with heterophily: survey. arXiv preprint arXiv:2202.07082, 2022a.
Wu, M., Zheng, X., Zhang, Q., Shen, X., Luo, X., Zhu, X., and Pan, S. Graph learning under distribution shifts: A comprehensive survey on domain adaptation, out- of-distribution, and continual learning. arXiv preprint arXiv:2402.16374 2024.
Zheng, X., Zhang, M., Chen, C., Li, C., Zhou, C., and Pan, S. Multi-relational graph neural architecture search with fine-grained message passing. In 2022 IEEE Interna- tional Conference on Data Mining (ICDM), pp- 783-792. IEEE, 2022b.
Wu, Q.. Zhang, H., Yan, J., and Wipf. D. Handling distri- bution shifts graphs: An invariance perspective. In International Conference Leurning Representations (ICLR), 2022.
Zheng, X., Liu, Y., Bao, Z.., Fang, M., Hu, X., Liew, A. W.-C., and Pan, S. Towards data-centric graph ma- chine learning: Review and outlook. arXiv preprint arXiv:2309.10977 2023a.
Xu, K., Hu, W., Leskovec, J., and Jegelka, S. How powerful are graph neural networks? In International Conference on Learning Representations, 2018.
Zheng, X., Zhang, M., Chen, C., Molaei, S., Zhou, C., and Pan, S. Gnnevaluator: Evaluating gnn performance on unseen graphs without labels. In Advances in Neural Information Processing Systems (NeurIPS), 2023b.
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-5823, 2020.
Zheng, X., Zhang, M., Chen, C., Nguyen, Q. H., Zhu, X., and Pan, S. Structure-free graph condensation: From large-scale graphs to condensed graph-free data. In Advances in Neural Information Processing Systems (NeurIPS), 2023c.
Yu, J., Liang, J., and He, R. Mind the label shift of augmentation- based graph ood generalization. In Pro- ceedings of the IEEE/CVF Conference on Computer Vi- sion and Pattern Recognition, pp. 11620-11630, 2023.
Yu, J., Zheng, Y., Koh, H. Y., Pan, S., Wang, T., and Wang, H. Collaborative expert ilms guided multi-objective molecular optimization. arXiv preprint arXiv:2503.03303, 2025.
10

---

Test-Time Graph Neural Dataset Search
Zheng, X., Zhang, M., Chen, C., Zhang, Q.. Zhou, C., and Pan, S. Auto-heg: Automated graph neural network heterophilic graphs. In Proceedings of the ACM Web Conference (WWW), pp. 611-620, 2023d.
In Advances in Neural Information Processing Systems (NeurIPS), 2022c.
Zheng, Y., Koh, H. Y., Yang, M., Li, L.. May, L. T., Webb, I., Pan, S., and Church, G. Large language models in drug discovery and development: From disease mecha- nisms clinical trials. arXiv preprint arXiv:2409.044p) 2024b.
Zheng, X., Song, D., Wen, Q., Du, B., and Pan, S. Online gnn evaluation under test-time bist-timega graph distribution shifts. In The Twelfth International Conference on Learning Representations, 2024a.
Zheng, Y., Koh, H. Y., Ju, J., Nguyen, A. T., May, I.. T. Webb, G. I., and Pan, S. Large language models for scien- tific discovery in molecular property prediction. Nature Machine Intelligence, pp. 1-11, 2025.
Zheng, Y., Pan, S., Lec, V. C., Zheng, Y., and Yu, P. S. Rethinking and scaling up graph contrastive learning: An extremely efficient approach with group discrimination.

---

Test-Time Graph Neural Dataset Search
Appendix
This is the appendix of our work: Test-Time Graph Neural Dataset Search With Generative Projection. Here, we provide additional details on the proposed PGNDS, including more related work discussions, further experimental settings such dataset statisties, and additional experimental results.
A. Related Work
Test-Time Adaptation. Test-time adaptation (TTA) focuses on dynamically adjusting pre-trained models to improve generalization on test samples (Jin et al., 2023; Wang et al., 2020; Chen al., 2023a; Liang et al., 2020; Jang et al., 2022), aiming at improving the performance and robustness of deep learning models in the deployment stage (Huang et al., 2021; Huang& & Chen, 2022). One of the earliest approaches, test-time training (TTT) (Sun et al., 2020), introduces self--uper learning with an auxiliary task, enabling model updates using single test sample. However, TTT requires access to training data to optimize the auxiliary task. To address this limitation, TENT (Wang et al., 2020) proposes fully test-time adaptation framework, which updates model parameters using only test data, making it more practical for real-world deployment.
Table AI. Comparison of different settings for different distribution shift related methods.
Test-time
Setting
Source Data
Target Data
Train Loss
Test Loss
Model Update
Fine-tuning
(10)
((x(),y(0)
Domain Generalization
(x)0) yra) (x(()) srlo (x((), (ro)) (xe(e) yloig
any (xit),
[(x(t), y((o)
Domain Adaptation
f(x()), y(1)) + [(x(0), x(()) + 2(x10,,x(0)
Unsupervised Domain Adaptation
x(1)
((x(9), (1))) + ((x((),
x
Test-time Training
x(1)
((x(t), y(d)) ) + ((x(mr)
((x11)
Test-time Adaptation (model-centric)
x(t)
((xit))
Test-time Adaptation (data-centric)
x(1)
[(x(())
Test-Time Adaptation on Graph. Despite the promising performance of TTA, existing methods predominantly model-centric and primarily designed for computer vision (Liang et al., 2024), making them less applicable to GNN models and graph data (Zhang et al., 2024d; Ju et al., 2024; Zhang et al., 2024c; Wu et al., 2024; Chen et al., 2022; Zheng et al., 2024a). Typically, TTA on graphs aims to dynamically fine-tune well-trained GNN models or modify test graphs to improve srrcctive-whethee adapting the model or the data data-existing methods can be classified into two categories: (a) test-time model adaptation (Chen et al., 2022; Wang et al., 2022; Zhang et al., 2024c;d); and (b) test-time graph adaptation (Jin et al., 2023).
GNN model generalization and inference performance. Based on the adaptation
Specifically, given the unseen test graphs, (a) test-time model adaptation mainly works on updating the well-trained GNN models using the self-supervised learning paradigm, where the primary objective is to optimize fine-tune the pre-trained GNN model parameters. For example, GT3 (Wang et al., 2022) and GraphTTA (Chen et al., 2022) employ self-supervised learning techniques to adapt GNN models at test time, targeting graph-level and node-level learning tasks, respectively. Moreover, НTmoTTT (Zhang et al., 2024d) proposes a homophily- -based and parameter-fie free graph contrastive learning task for fully test-time GNNs training.
In contrast, (b) test-time graph adaptation takes a data-centric approach, modifying the test graph data while keeping the pre-trained GNN model parameters unchanged. Typically, GTRANS (Jin al., 2023) modifies test graph data without accessing the training procedure or GNN architectures. However, GTRANS relies on fully parameterized matrix to adjust test-time node features, which restricts its ability to effectively model updated test graphs. Additionally, its binary-space projected gradient descent limits flexibility in handling diverse graph structures. GraphPatcher (Ju et al., 2024) addresses the test-time graph degree shift problem by repatching graphs for node-level tasks. However, its approach is limited to node degree-based modifications, which restricts its flexibility in adapting to diverse test-time distribution shifts the graph level.
12

---

Test-Time Graph Neural Dataset Search
Table A2. Dataset statistics for graph classification and regression on protein and molecular graphs.
# Nodes
Edges
Graph Types
Datasets
Task Types
# Graphs
# Tasks
# Train / Test
Avg/Max
Avg./Ma
Protein
Enzymes
Classification
587
33.0/ 125
63.2 149
470 117
Ogbg-BACE
Classification
1,513
34.1 /
73.7 / 202
1210 152
Oghg-BBBP
Classification
2,039
24.1 / 132
51.9/ / 290
1631 204
Molecular
Ogbg-Cline Ogbg-ClinTox
Classification
1,477
26.2 136
55.8 286
1181/ 148
Oghg-FreeSolv
Regression
642
8.7
16.8 50
513 65
QM9
Regression
133,885
8.8/9
9.4 /13
120803 13082
In summary, our work addresses the broader challenges of flexibility and generalization in test-time adaptation for GNNs by learning a parameterized test-time graph distribution. This enables improved inference performance on unseen test graphs using well-trained GNNs. Unlike existing graph data-centric TTA methods, our approach captures wider range of distribution shifts and is particularly effective handling unknown distribution shifts the graph level. Furthermore, it is important highlight why certain baselines GraphTTA, and HomoTTT, lack publicly available code, making replication infeasible. Others, like GraphPatcher, operate in fundamentally different learning scenarios that are not directly comparable to our setting. Given these constraints, we carefully selected GTRANS, the most recent state-of-the-art method, as the primary baseline for fair and meaningful comparison. This ensures that our evaluation remains rigorous and aligned with the intended scope of test-time graph adaptation.
not included in our experiments. Some methods, such GT3,
Distribution Shift Related Methods. We provide a comparison of different methods addressing distribution shift problems across various learning settings. It outlines the key characteristics for each setting based the following aspects as shown in Table A1:
Source Data: The availability of labeled source data during training.
Target Data: The type of target data used during testing.
Train Loss: The objective functions used during training, indicating whether source data, target data, or both are used.
Test Loss: The loss function evaluated during test time, if applicable.
Test-time Model Update: Whether the model undergoes parameter updates during test time (V for yes, x for no).
The table highlights distinctions between:
Common learning paradigms like fine-tuning, domain generalization, domain adaptation, and unsupervised domain adaptation.
Test-time training approaches, which incorporate test-time loss evaluation.
Test-time adaptation, divided into model-centric (updating model parameters) and data-centric (focusing adapting data rather than modifying model parameters).
The comparison underscores the focus of test-time adaptation methods on improving generalization without relying on access to source data during testing. In this work, we focus primarily on test-time graph adaptation using data-centric methods. Given that different distribution shift related approaches suit varying settings, we excluded comparison methods that pertain to distinct application scenarios.
B. Dataset Detail
We provide an overview of dataset statistics used for graph classification and regression tasks on protein and molecular graphs in Table A2. The table highlights the diversity of datasets regarding scale, task type, and structural complexity, showeasing their suitability for evaluating test-time graph adaptation methods for graph-level tasks.
13

---

Test-Time Graph Neural Dataset Search
C. Hyper-parameter A Analysis
We present more hyper-parameter sensitivity analysis for the Ogbg-BBBP dataset on graph classification, measured using ROC-AUC in Fig. A1. The analysis explores the impact of paper-pprameterss a, 3, and 7y in the dual conditional diffusion framework when B 1 and y
In general, we can observe that the performance (ROC-AUC) is visualized for varying a, B y values. And moderate values of o, B, and are critical for optimal performance, when extreme values tend to degrade the model's effectiveness.
D. Visualization Comparison
Fig. A2 illustrates a visualization comparison between the original test graph and the proposed PGNDS -eWDDS-geeeetted test graph for molecules in the QM9 dataset for property The comparison demonstrates that while key features such as atom connectivity and bond types are preserved, the PGNDS-generinn versions exhibit notable differences in graph structure, highlighting the effectiveness of our PGNDS for adapting molecular graphs while maintaining their fundamental chemical properties.
Hyperparameter Analysis (B=1) 1)
Hyperparameter Analysis (y=1)
0.01
0.70
0.01
0.87
0.65
0.63
0.54
0.54
0.54
0.64
0.64
0.55
0.51
0.51
0.63
0.69
0,1
0.68
0.67
0.08
0.09
0.08
0.7
0.64
0.2
0.69
0.66
0.68
0.68
0.7
0.71
0.77
0.87
0.5
0.69
0.7
0.7
0.66
0.67
0.68
0.69
0.66
0.68
0.7
0.7
0.69
0.64
0.69
0.01
0.2
0.01
0.1
0.2
0.5
5
B
Figure AI. Slyper-paramete sensitivity analysis on Ogbg-BBBP for graph classification (ROC-AUC) with a. G. and yy in Eq. (12) for different constrains dual conditional diffusion.
QM9-A original (G-A.)
QM9-A PGNDS (G.9).
(e) QM9-a original (G!21)
(d) QM9-A PGNDS (G.21)
Figure A2, Visualization comparison between original test graph and our proposed PGNDS generated test graph for molecules in QM9.
14


Reviews:
[{"review_id": "0WOLovBzSN", "reviewer": "Reviewer_G3YQ", "summary": "The authors introduce a new problem, test-time graph neural dataset search, to learn the optimal distribution of unknown test graph datasets. For this purpose, they propose PGNDS, a generative projection driven by a diffusion model. By projecting test graphs back to the training distribution, PGNDS learns test-time adaptation by generating refined test graphs. Experimental results demonstrate the effectiveness of PGNDS.", "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": "**[Re-Claims and Evidence (1): Difference between “Test-Time Graph Neural Dataset Search” and “Test-Time Graph Adaptation”]**\n\n>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.\n\n>| Aspect                         | Test-Time Graph Adaptation                         | Test-Time Graph Neural Dataset Search      |\n>|-------------------------------|----------------------------------------------------|------------------------------------------------------|\n>|Granularity               | Per-graph (adapts each test graph individually)    | Dataset-level (searches a distribution over graphs)  |\n>|Adaptation Objective      | Modify a single test graph to improve prediction locally  | Generate task-aligned graph distributions toward training domain  |\n>| Method Type               | Direct feature and structure manipulation | Generative modeling + sampling via diffusion         |\n>| Optimization Target       | One graph at a time                                | Entire test-time graph set as a distribution         |\n\n---\n**[Re-Claims And Evidence (2)] Whether test-time distribution projection change the properties of refined test graphs?**\n\n> 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.\n\n>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.\n\n---\n\n**[Re-Methods and Evaluation Criteria] Applicability to Larger Graphs**\n>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.\n\n---\n**[Re-Experimental Designs or Analyses] Ensemble Inference**\n>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.\n\n---\n\n**[Re-Weakness and Suggestions] Recent methods.**\n>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.\n\n---\n\n**[Re-Questions For Authors] On Figure A2**\n\n>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.", "claims_and_evidence": "What is the specific difference between the problem “test-time graph neural dataset search” and “test-time graph adaptation”? \n\n\nWould projecting the entire test graph set distribution back to the training set change the properties of refined test graphs?", "methods_and_evaluation_criteria": "This method is applicable to small molecular and protein graphs. Is it applicable to larger graphs?", "theoretical_claims": "The theoretical analyses seem to be correct.", "experimental_designs_or_analyses": "The “ensemble inference” phase aggregates the representative information from both the original and adapted test graphs. Experiments should their individual results, and explain why the method does not use only the adapted graph. Because the adapted graph seems to be the best structure.", "relation_to_broader_scientific_literature": "Compared with the test-time model adaptation methods, this paper avoids modifying GNN parameters and is more applicable. This paper also proposes an innovative method.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths: The proposed method is novel for test-time graph-level tasks.\n\nWeaknesses: Recent methods are not compared in the experiment.", "other_comments_or_suggestions": "This paper should compare more recent methods, including test-time model adaptation methods and graph adaptation methods they refer to.", "questions_for_authors": "In figure A2, I cannot understand the effect of graph structure adaption. Why (b) and (d) is better than (a) and (c)?", "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": {"ethical_review_concerns": "No"}}, {"review_id": "kG6KuS307J", "reviewer": "Reviewer_VWGF", "summary": "This paper introduces test-time graph neural dataset search with generative projection to improve test-time adaptation for Graph Neural Networks (GNNs) facing distribution shifts. The proposed method, PGNDS, uses a generative projection approach to refine test graphs without modifying the trained GNN. PGNDS consists of three key steps: dual conditional diffusion for graph generation, dynamic search for selecting the best test graphs, and ensemble inference for improved predictions. Experiments on real-world graph datasets shows better performance over baselines.", "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": "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.\n\n---\n\n**[Re-Weakness-(1)] On Method Built on Multiple Concepts**\n>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:\n\n> - Graph diffusion: a generative process that models test-time distribution shifts via reverse-time diffusion;\n> - Conditional guidance: a mechanism to control the reverse-time diffusion for generating refined test graphs using pretrained GNNs' knowledge;\n> - Dataset-level search: selecting or generating a refined set of test graphs by optimizing over the entire test distribution.\n\n>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.\n\n---\n\n**[Re-Weakness-(2)] On Relation to Graph Neural Architecture Search (GNAS)**\n>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.\n\n---\n\n**[Re-Weakness-(3)] On Dynamic Search and Stopping Criterion**\n>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.\n\n---\n**[Re-Questions For Authors] On Generalization to Other Graph Types and Tasks**\n>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.", "claims_and_evidence": "Yes,  no problematic claims are found.", "methods_and_evaluation_criteria": "Yes, make sense.", "theoretical_claims": "Did not go into much detail, but generally checked.", "experimental_designs_or_analyses": "Generally checked.", "relation_to_broader_scientific_literature": "Relatively related to existing graph learning problem and GNNs.", "essential_references_not_discussed": "N/A.", "other_strengths_and_weaknesses": "For strengths:\n\n(1)The paper introduces a new learning problem for handling test-time adaptation in graph data, which appears to be a novel challenge for GNN generalization at inference time. The idea of adapting test graphs without modifying model parameters seems practical for real-world applications.\n\n(2)The overall writing is organized well, the background about test-time adaptation on graphs and challenges is described relatively clearly. Each module of PGNDS method, i.e., leveraging dual conditional diffusion, dynamic search, and ensemble inference, are well-designed and well-described, although it is a little complex to understand.\n\n(3)This paper includes several experiments that show the proposed method performs better than existing approaches, showing the proposed PGNDS could capture unseen test graph distributions through a generative method guided by well-trained GNNs.\n\nFor weaknesses:\n\n(1)This method involves many concepts like graph diffusion, conditional guidance, and dataset search, it is quite difficult to fully understand it.\n\n(2)How does this work relate to graph neural architecture search? should this be discussed?\n\n(3)Some parts, like how the \"dynamic search\" process selects useful graphs, are not fully explained, more details need to be explained further. How to stop the search and guarantee the selected test graphs are optimal?", "other_comments_or_suggestions": "See the previous question.", "questions_for_authors": "See the weakness. \n\nAnother question is about the generalization ability of the model. The evaluation primarily focuses on molecular graphs for graph classification and regression tasks. Would the proposed PGNDS be equally effective on other types of graphs, such as social networks or citation graphs, or on different graph learning tasks, like node classification? If so, could you explain how PGNDS would adapt to these different graph types and tasks?", "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": "7JLwo5lPG1", "reviewer": "Reviewer_TmwA", "summary": "This paper tackles a really interesting and new problem called 'test-time graph neural dataset search.' It enables GNN models handle data they’ve never seen before at test time by creating new graphs similar to the training set. To tackle this problem, the authors propose PGNDS, a method that reconstructs an unseen test graph distribution by leveraging a well-trained GNN as a guiding mechanism. The proposed PGNDS framework contains graph diffusion model as a generator, for generating new test distributions through test-back-to-training distribution mapping, and it uses dynamic search to select new test graphs, and uses new generated graphs and original test graphs to provide final predictions. This work reports the main results and other auxiliary experiments on real-world test graphs.", "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": "**[Re-Weakness(1)] Computation of Dual Conditional Diffusion and Dynamic Search**\n>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.\n\n---\n**[Re-Weakness(2)] Well-trained GNNs and Diffusion Model**\n>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.\n\n---\n\n**[Re-Other Comments Or Suggestions] Limitations and Future Directions**\n>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).\n\n>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.\n\n---\n\n**[Re-Questions(1)] Under What Circumstances the Generative Projection Process Might Fail**\n\n>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.\n\n---\n\n**[Re-Questions(2)] Ensemble Inference Ratio**\n>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.\n\n---\n\n**[Re-Questions(3)] Distribution Alignment and Generalization Guarantee**\n>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.\n\n>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.", "claims_and_evidence": "The authors clearly state their claims, and personally, I find the evidence convincing overall, though some practical implications could be explained more clearly.", "methods_and_evaluation_criteria": "I find the proposed three-stage PGNDS framework interesting and quite practical. Instead of tuning the model itself, it cleverly changes the graph data at test time. The design of PGNDS generally aligns well with the proposed test-time graph neural dataset search problem. The evaluation criteria (ROC-AUC, RMSE) used benchmark datasets (molecular and protein), and baselines, on graph learning tasks are appropriate.", "theoretical_claims": "I quickly checked the theoretical parts, and they seem logical. However, it would help if the authors could simplify or clarify the intuition behind these theories.", "experimental_designs_or_analyses": "I briefly checked the soundness/validity of the experimental designs and analyses, as well as the ablation study, hyperparameter analysis, and running time comparison. The experiments look good overall, no issues need to be discussed here.", "relation_to_broader_scientific_literature": "This work broadly discusses its connection to the graph distribution shift problem in the supplementary material, particularly in “Table A1, Comparison of different settings for different distribution shift related methods”. The relationship between this work and other research methods is relatively clear and helps distinguish it from existing approaches.", "essential_references_not_discussed": "Not sure, maybe other graph diffusion/generation methods.", "other_strengths_and_weaknesses": "### Key strengths:\n- [Originality] Introduces test-time graph neural dataset search - a novel extension of TTA using generative projection. Shifts focus from graph-level adaptation to distribution-level optimization.\n- [Significance] Addresses critical GNN deployment challenge: performance degradation from graph distribution shifts. Practical solution through training distribution projection enhances real-world robustness.-\n- [Technical Clarity] Well-structured modular design (dual diffusion, dynamic search, ensemble inference) with rigorous mathematical formulation and theoretical grounding.\n- [Validation] Comprehensive experiments across datasets show PGNDS outperforms SOTA baselines. Ablation studies validate component contributions and parameter sensitivity.\n### Weaknesses:\n- [Complexity] While the proposed generative projection method is innovative, the dual conditional diffusion and dynamic search procedures may introduce significant computational overhead during inference, potentially limiting its applicability in resource-constrained environments.\n- [Dependence on Well-trained GNN Models] The effectiveness of PGNDS heavily depends on well-trained pre-existing GNN and diffusion models. This assumes extensive initial training data and reliable model training, whether these assumptions are rational?", "other_comments_or_suggestions": "It would be beneficial to include a more detailed discussion on the limitations of the proposed PGNDS and future potential directions of the proposed test-time graph neural dataset search problem.", "questions_for_authors": "Here are my questions:\n\n1.The generative projection process seems involves randomness through graph diffusion model, then, under what circumstances might the generative projection process fail or lead to degraded performance?  are there any potential theoretical or practical limitations?\n\n2.The paper introduces ensemble inference as a final step. How is the optimal ratio between original test graphs and newly generated graphs determined for achieving the best inference results?\n\n3.The paper claims that the test-time graph distribution is mapped back to the training distribution using a diffusion model. However, is there a formal guarantee that the generated graphs truly align with the training distribution in a way that improves generalization, rather than simply generating more training-like graphs that may not be representative of the test distribution?", "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": {"ethical_review_concerns": "N/A"}}, {"review_id": "HbuEtyvCZ2", "reviewer": "Reviewer_AMWj", "summary": "This paper addresses the test-time adaptation challenge in graph neural networks (GNNs). \nThe main focus is on the generalization about graph data and test-time GNN inference. \nThe main challenge is that GNN models trained on training graphs may not work well on a new, unseen test graph. The authors propose a new learning problem, test-time graph neural dataset search, along with a generative Projection based test-time Graph Neural Dataset Search method, named PGNDS, to adjust the model at test time so it can better handle new test graphs. PGNDS uses a graph diffusion generator to map new generated test graph to a new distribution close to the training distribution and this process is guided by well-trained GNNs with different constraints for the reverse diffusion process. \nThe main results of the proposed PGNDS on real-world datasets, show that PGNDS could improve the GNN test time inference results.", "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": "**[Re-Weakness(1)] Computation and Explanations of Multiple Module Components in PGNDS**\n>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:\n\n> - Dual conditional diffusion: a generative process that models test-time distribution shifts via reverse-time diffusion.\n> - Dynamic search: selecting or generating a refined set of test graphs by optimizing over the entire test distribution.\n> - Ensemble inference: aggregating information from the original test graphs and refined adapted graphs to obtain optimal test-time inference performance.\n\n>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.\n\n>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.\n\n---\n\n**[Re-Weakness(2)] Feasibility of Well-Trained GNN and Diffusion Models**\n>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.\n\n---\n\n**[Re-Questions(1) & (2)] Code and Intuitive Explanation of Dual Conditional Diffusion and Guidance Terms**\n>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.\n\n>For an intuitive explanation of dual conditional diffusion and the multiple guidance terms:\n>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:\n\n> - 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.\n> - 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.\n> - Third, graph diversity (**$r_{\\text{gdist}}$**) prevents the process from collapsing into trivial solutions, maintaining variety in the generated test-time graph candidates.\n\n>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].\n\n---\n\n**[Re-Questions (3) ] Robustness to Adversarial Perturbations**\n>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.\n\n>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.", "claims_and_evidence": "Yes, the overall paper is relatively clear, providing detailed descriptions of the proposed method and theoretical justifications.", "methods_and_evaluation_criteria": "Yes, the proposed PGNDS methods and evaluation criteria (six molecular and protein datasets) are reasonable for the problem of test time GNN adaptation and the practical application.", "theoretical_claims": "Yes, the theoretical claims specifically in Section 2.4. Theoretical Justification are checked.", "experimental_designs_or_analyses": "Yes, the experimental design is appropriate overall, and the experimental results could verify its effectiveness.", "relation_to_broader_scientific_literature": "The key contribution of this paper is to formulate the test-time graph adaptation problem to a new test-time graph dataset generation problem, different from existing test-time adaptation methods on graph, it moves beyond direct model tuning and graph nodes/structure modification.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "The paper demonstrates several notable strengths. \n\nFirst, it innovatively propose new test‐time adaptation problem by proposing “graph neural dataset search”. This creative formulation leverages generative projection and dual conditional diffusion to remap unseen test graph distributions toward the training distribution, which is a fresh perspective that could inspire further work in test‐time adaptation for graphs. Second, the modular design (dual conditional diffusion, dynamic search, ensemble inference) is well-structured, and the inclusion of theoretical justifications adds depth. Last, the extensive experimental evaluation—with thorough ablation study and hyperparameter sensitivity analyses—also validate the performance of the proposed PGNDS method on molecular and protein graph tasks. \n\nOn the other hand, some potential weaknesses deserve mention. \n\nFirst, integrating multiple modules (dual conditional diffusion, dynamic search, and ensemble inference) that may impose a high computational overhead and could be challenging to reproduce. Simplifying some aspects or providing additional intuitive explanations might be helpful. Second, the approach assumes access to well-trained GNN and diffusion models, which may not always be feasible.", "other_comments_or_suggestions": "No.", "questions_for_authors": "First, given the complexity of the overall framework, will code be release to ensure the reproducibility of the results?\n\nSecond, the theoretical justification involving dual conditional diffusion and multiple guidance terms is quite dense. Could you offer a more intuitive explanation of how these components interact in practice to effectively reduce test-time distribution shifts?\n\nLast, how the proposed PGNDS would perform against adversarial perturbations or noisy test-time graphs? Understanding its resilience to noisy data would be valuable for real-world applications.", "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: This paper introduces a novel and timely framework, PGNDS, which tackles the problem of test-time adaptation for Graph Neural Networks (GNNs) through a new paradigm termed test-time graph neural dataset search. Instead of adapting models or individual graphs, PGNDS performs distribution-level test graph refinement using a generative diffusion process guided by pretrained GNNs. The framework is modular and theoretically grounded, incorporating dual conditional diffusion, dynamic search, and ensemble inference. Reviewers largely agree on the significance and originality of the problem formulation, the technical soundness of the proposed solution, and the comprehensiveness of the experimental validation. While some reviewers noted the complexity of the framework and the need for clearer explanations or expanded evaluation, the authors provided detailed and thoughtful rebuttals that addressed these points effectively. Given its originality, well-executed design, and strong empirical performance, I recommend acceptance.
Decision: Accept (poster)