Test-Time Graph Rebirth: Serving GNN Generalization Under Distribution Shifts

Authors: Xin Zheng, Yu Zheng, Qin Zhang, Haishuai Wang, Yuxuan Liang, Alan Wee-Chung Liew, Shirui Pan
Primary area: learning on graphs and other geometries & topologies
Keywords: graph neural networks, distribution shifts, test-time inference, MLOps

Abstract:
Distribution shifts between training and test graphs typically lead to the decreased performance of graph neural networks (GNNs) with suboptimal generalization in real-world applications. Despite advances in graph learning under distribution shifts through designing various model architecture development with customized training strategies, existing solutions can be challenging in practical GNN deployment because they often require significant modifications or retraining of the GNNs. To address such challenges, in this work, we propose a novel method, i.e., Test-Time Graph REBirth, dubbed TT-GREB, to effectively generalize the well-trained GNN models to the test-time graphs under distribution shifts by directly manipulating the test graph data. Concretely, we develop an overall framework designed by two principles, corresponding to two submodules: (1) prototype extractor for re-extracting the environment-invariant features of the test-time graph; and (2) environment refiner for re-fining the environment-varying features to explore the potential shifts. Furthermore, we propose a dual test-time graph contrastive learning objective with an effective iterative optimization strategy to obtain optimal prototype components and environmental components of the test graph. By reassembling these two components, we could obtain a newly reborn test graph, which is better suited for generalization on the well-trained GNN model with shifts in graph distribution. Extensive experiments on real-world graphs under diverse test-time distribution shifts could verify the effectiveness of the proposed method, showcasing its superior ability to manipulate test-time graphs for better GNN generalization ability.

Decision: reject
Reviewer scores: 3, 5, 5, 5

Reviews and discussion:
<thread>
creation date:
1729755080447

thread_type:
official_review

<thread_object>
creator: reviewer_aceo

<content>
summary:
In this paper,  the authors introduce a novel graph data-centric paradigm for enhancing the ability of well-trained GNNs to real-world graphs experiencing distribution shifts at test time. The uniqueness of this paper lies in the re-operation of the test data to model the invariant features, which are then fed into the trained GNN for prediction. Empirically, experiments are conducted on five datasets and four benchmark networks to verify the effectiveness of the model.

strengths:
1. The research question is critical
2. It is interesting to design a separate model to capture invariant representations on test data
3. Experimental verification was performed on multiple backbone networks

weaknesses:
1. This method can be seen as the process of designing an unsupervised model to extract invariant features during the test phase. However, it is often difficult to perform decoupled learning under unsupervised conditions.
2. This is misleading in Figure 1(a). In the general graph OOD generalization setting, the test graph is not seen during training. The figure may be a problem setting of domain adaptation.
3. This approach does not seem to be limited to graph data. What are the unique challenges for graph data?
4. In the experiment, the comparison method is not new enough, the latest one comes from ICLR2023. At the same time, Figure 5 is a bit difficult to observe, and its readability can be further improved. If new comparative experiments can be provided here, I will consider improving my score.

questions:
Please see the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730625938005

thread_type:
official_review

<thread_object>
creator: reviewer_ooog

<content>
summary:
This paper introduces a data-centric approach called TT-GREB to address the distribution shift in graph learning at test time. Specifically, TT-GREB employs 1. a prototype-extractor module to extract distribution-invariant subgraph, and 2. an environment-refiner module to extract spurious subgraph. The distribution-invariant subgraph, spurious subgraph, and the new graph of the previous two construct a triplet to supervise the graph neural network at test time to mitigate the distribution shift. The prototype-extractor module, environment-refiner module, and the graph neural network are optimized iteratively together. Empirical studies on several datasets show that the proposed TT-GREB can improve the out-of-distribution generalization performance of graph neural network.

strengths:
Improving graph neural networks (GNNs)' out-of-distribution (OOD) performance is important to their real-world deployment since different distribution shifts may occur in the real world. Although some distribution-invariant GNN architectures have been designed, it is still unrealistic to wish these models to generalize to arbitrary unseen distributions. Adapting GNNs at test time is more realistic and general as it can employ some hints from testing data to guide network generalization. This paper follows the test-time adaptation approach and the proposed method can work well on several datasets.

weaknesses:
Yet, the reviewer is concerned about some technical details as follows:

1. Why use two modules to extract the distribution-invariant and spurious subgraphs respectively? According to Proposition 1, the distribution-invariant subgraph and spurious subgraph are fully complementary. Thus, using one module to extract any one of the two subgraphs is sufficient to obtain both distribution-invariant and spurious subgraphs. Instead, this paper introduces two modules to extract these two subgraphs.

2. Based on 1, why is G_{te}^{'} different from G_{te}? As G_{te}^{'} is constructed by the distribution-invariant and spurious subgraphs from G_{te},  G_{te}^{'} should be the same as G_{te} according to 1.

3. As G_{te}^{'} has a spurious subgraph, why push its embedding closer to the embedding of the distribution-invariant subgraph?

4. EERM is designed to improve the OOD generalization ability of GNNs, why does it underperform most baselines in Table 1?

questions:
The authors are encouraged to address the concerns in the Weaknesses part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730671883919

thread_type:
official_review

<thread_object>
creator: reviewer_hkp5

<content>
summary:
This paper introduces a test-time data-centric approach, Test-Time Graph REBirth (TT-GREB), designed to improve the generalization abilities of pre-trained Graph Neural Networks (GNNs) under graph distribution shifts at test time. It presents a method for regenerating test-time graphs by separating environment-invariant features (which capture the core predictive patterns) from environment-varying features (which address distributional shifts).  TT-GREB leverages dual test-time graph contrastive learning objective with self-supervision signals, along with an effective iterative optimization strategy to obtain expressive prototype features and environmental features. Extensive evaluations demonstrate TT-GREB's superiority over state-of-the-art methods in enhancing GNN performance under test-time distribution shifts.

strengths:
1. The approach seems novel to me. This paper presents an approach to test-time graph rebirth, which contrasts with existing methods by focusing on modifying architectures and training strategies to address distribution shifts. The motivation is also clear to me.
2. The idea to extract environment-invariant features and refine the environment-varying features is reasonable to me.
3. Experiments include various GNN backbones and baselines to compare the performance.

weaknesses:
1. The experimental results are comparable to the data-centric baseline GTRANS. The improvement is not very evident.
2. The proposed method appears less effective on large-scale datasets, such as the arXiv dataset. For even larger-scale datasets, will the performance be even worse and will the iterative optimization strategy be computationally intensive when the graph sizes are very large?
3. The captions and ticks in figures are too small for readers to read.
4. Some figures have inconsistent font sizes in their captions (e.g., Figure 4), detracting from the professional presentation.

questions:
1. Can the method extend to other model architectures like Graph Transformers?
2. In terms of runtime comparison, could the authors provide more detailed comparisons of training and inference times?
3. How well does TT-GREB generalize to larger-scale datasets beyond those used in the paper? Will it hurt the performance or lead to huge time complexity? Could the authors clarify the scalability of the iterative optimization strategy, especially when applied to very large graph datasets?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699455609

thread_type:
official_review

<thread_object>
creator: reviewer_msyu

<content>
summary:
The paper presents a novel data-centric approach called Test-Time Graph Rebirth (TT-GREB) to improve the generalization of GNNs under distribution shifts. The method manipulates test graph data directly and does not require modifications to the GNN architecture or retraining. To achieve this, the method decomposes the test graph into environment-invariant and environment-varying components, using a prototype extractor and environment refiner. This leads to the creation of a “reborn” test graph that better matches the training graph distribution, enhancing the GNN’s generalization without altering its original structure. The dual contrastive learning objective and iterative optimization are key to refining the test graph for improved node classification performance.

strengths:
The paper introduces a practical and innovative solution to the problem of GNN generalization under test-time distribution shifts without requiring retraining or modifications to the GNN model itself. The dual contrastive learning framework is designed to refine test-time graphs by optimizing environment-invariant and environment-varying features. Experimental results on diverse real-world datasets demonstrate that TT-GREB outperforms existing methods.

weaknesses:
1. The problem setup seems quite contrived. Particularly, it is mostly unclear why a graph could be explicitly decomposed into environment-invariant and environment-varying parts, where the environment-varying variables are upper bounded with a certain notion of discrepancy (which is also not formally defined). While the authors write this as a proposition, it seems more like an (unrealistic) assumption without any support or proof.
2. There is no guarantee that the proposed Test-time Graph Rebirth as in definition 3.1 can improve out-of-distribution generalization.
3. The method only shows marginal improvement over another data-centric graph ood method.

questions:
See above
</content>
</thread_object>

</thread>
<thread>
creation date:
1732232975288

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common Response to All Reviewers

comment:
We sincerely thank all the reviewers for their constructive feedback, insightful comments, and recognition of the strengths of our work. We are particularly grateful for the recognition of:  

- *Important/critical research problem* of our proposed Test-time Graph REBirth (TT-GREB) to real-world deployment (**Reviewer oooG, AcEo**)
-  *Novel/practical and innovative/uniqueness/reasonable/interesting/more realistic and general/ approach* of graph data-centric paradigm for test-time data modification, under test-time distribution shifts (**All Reviewers!**)
-  *Effective iterative optimization strategy* with optimized *environment-invariant and environment-varying features* (**Reviewer hkp5, Reviewer mSYU**)
- *Outperforming* experimental results on *multiple datasets and GNN backbones* over state-of-the-art methods  (**All Reviewers!**)
- Clear motivation (**Reviewer hkp5**)

We are grateful for your positive recognition and have carefully addressed your insightful concerns in the detailed responses below. If our explanations and additional evidence resolve your questions, we would greatly appreciate it if you could reconsider your evaluation of our work. Thank you again for your thoughtful reviews and valuable feedback.
</content>
</thread_object>

</thread>


Metareview:
metareview:
(a) The paper proposes a method called Test-Time Graph REBirth (TT-GREB) to improve the generalization of graph neural networks (GNNs) under distribution shifts. It aims to address the issue by directly manipulating test graph data through a prototype extractor and an environment refiner. The key findings are that TT-GREB can outperform existing methods in enhancing GNN performance on diverse real-world datasets under test-time distribution shifts.

(b) The strengths of the paper include proposing a new data-centric approach for test-time graph modification, introducing a dual contrastive learning framework for optimizing test-time graphs, and demonstrating effectiveness through extensive experiments with various GNN backbones and datasets.

(c) The weaknesses are that the problem setup seems contrived with unclear assumptions (e.g., graph decomposition), the improvement over a data-centric baseline is not highly significant, and the method may face scalability issues on large-scale datasets with high computational intensity.

(d) The main reasons for rejection are the lack of clear justification for the proposed approach, especially regarding the assumptions. Although the experimental results show some improvements, they are not overwhelmingly convincing, and the concerns about scalability and the overall problem formulation remain unresolved.