[Paper Review] Test-Time Training for Graph Neural Networks

Yiqi Wang, Chaozhuo Li+5 more|arXiv (Cornell University)|Oct 17, 2022

Advanced Graph Neural Networks4 citations

TL;DR

This paper proposes GT3, the first test-time training framework for Graph Neural Networks (GNNs) to improve generalization under distribution shifts in graph classification. By applying self-supervised learning and an adaptation constraint during inference, GT3 fine-tunes the GNN on each test graph, significantly reducing performance gaps between training and test sets, especially on datasets like ogbg-molhiv and D&D with large distribution shifts.

ABSTRACT

Graph Neural Networks (GNNs) have made tremendous progress in the graph classification task. However, a performance gap between the training set and the test set has often been noticed. To bridge such gap, in this work we introduce the first test-time training framework for GNNs to enhance the model generalization capacity for the graph classification task. In particular, we design a novel test-time training strategy with self-supervised learning to adjust the GNN model for each test graph sample. Experiments on the benchmark datasets have demonstrated the effectiveness of the proposed framework, especially when there are distribution shifts between training set and test set. We have also conducted exploratory studies and theoretical analysis to gain deeper understandings on the rationality of the design of the proposed graph test time training framework (GT3).

Motivation & Objective

• To address the persistent performance gap between training and test sets in graph classification, especially under distribution shifts.
• To design a test-time training framework that adapts pre-trained GNNs to individual test samples without retraining from scratch.
• To ensure robustness against feature space distortion during test-time adaptation by incorporating a constraint mechanism.
• To validate the effectiveness and rationality of test-time training in the graph domain through empirical and theoretical analysis.

Proposed method

• GT3 employs a self-supervised learning (SSL) task that leverages both local and global graph structure and node attributes to generate supervisory signals for test-time adaptation.
• The framework introduces an adaptation constraint that regularizes the model updates during test-time training to prevent harmful feature space distortion.
• It uses shared layers between the main task and SSL task, with the number of shared layers tuned to balance generalization and adaptation capacity.
• The model performs iterative updates on each test graph using a contrastive learning objective that compares augmented views of the same graph.
• The method is applied during inference only, with no changes to the pre-trained GNN weights during training.
• Ablation studies confirm the necessity of both the SSL task and the adaptation constraint for optimal performance.

Experimental results

Research questions

• RQ1 Can test-time training effectively reduce the performance gap between training and test sets in graph classification under distribution shifts?
• RQ2 How can self-supervised learning be effectively designed for graph data to guide test-time adaptation?
• RQ3 What role does the adaptation constraint play in preventing feature space distortion during test-time fine-tuning?
• RQ4 How do the number of shared layers and the choice of SSL components affect the overall performance of GT3?
• RQ5 Is test-time training for GNNs theoretically and empirically justified for improving model generalization?

Key findings

• GT3 significantly improves graph classification performance on benchmark datasets such as ogbg-molhiv, D&D, ENZYMES, and PROTEINS, especially under distribution shifts.
• On ogbg-molhiv, GT3 reduces the performance gap between training and test sets, with GIN-based models achieving up to 75.2% ROC-AUC on test, compared to 84.1% on validation.
• Abllation studies show that removing the adaptation constraint, global contrastive learning, or local contrastive learning degrades performance, proving the necessity of all components.
• The number of shared layers influences performance, with optimal configurations boosting accuracy—e.g., on D&D, shared layers improve performance over baseline.
• Theoretical and empirical analysis confirms that test-time training with SSL is rational and effective for GNNs, particularly in out-of-distribution settings.
• GT3 demonstrates consistent gains across multiple datasets and GNN architectures, validating its generalization capability.