Test-Time Training on Graphs with Large Language Models (LLMs)

Jiaxin Zhang,Yiqi Wang,Xihong Yang,Siwei Wang,Yu Feng,Yu Shi,Ruicaho Ren,En Zhu,Xinwang Liu

2024-04-21

Abstract:Graph Neural Networks have demonstrated great success in various fields of multimedia. However, the distribution shift between the training and test data challenges the effectiveness of GNNs. To mitigate this challenge, Test-Time Training (TTT) has been proposed as a promising approach. Traditional TTT methods require a demanding unsupervised training strategy to capture the information from test to benefit the main task. Inspired by the great annotation ability of Large Language Models (LLMs) on Text-Attributed Graphs (TAGs), we propose to enhance the test-time training on graphs with LLMs as annotators. In this paper, we design a novel Test-Time Training pipeline, LLMTTT, which conducts the test-time adaptation under the annotations by LLMs on a carefully-selected node set. Specifically, LLMTTT introduces a hybrid active node selection strategy that considers not only node diversity and representativeness, but also prediction signals from the pre-trained model. Given annotations from LLMs, a two-stage training strategy is designed to tailor the test-time model with the limited and noisy labels. A theoretical analysis ensures the validity of our method and extensive experiments demonstrate that the proposed LLMTTT can achieve a significant performance improvement compared to existing Out-of-Distribution (OOD) generalization methods.

Machine Learning,Artificial Intelligence

What problem does this paper attempt to address?

This paper proposes a solution to the performance degradation of Graph Neural Networks (GNNs) when there is a discrepancy between the training data and the testing data distribution, namely Out-of-Distribution (OOD). Traditional Test-Time Training (TTT) methods rely on unsupervised strategies to adapt to the testing data. However, inspired by the annotation capability of Large Language Models (LLMs) on Text-Attributed Graphs (TAGs), the authors propose a new method called LLMTTT. LLMTTT utilizes the annotation of LLMs on a carefully selected set of nodes for testing time adaptation and adopts a hybrid active node selection strategy that considers node diversity, representativeness, and the prediction signal of pre-trained models. To deal with the limited and noisy labels generated by LLMs, a two-stage training strategy is designed to adjust the model.

The main contributions of this paper include:
1. Introducing a new pipeline LLMTTT that utilizes LLMs as annotators to obtain pseudo-labels for test-time training to alleviate the OOD problem on graph data.
2. Developing a hybrid active node selection strategy that considers not only node diversity and representativeness, but also the prediction signal of pre-trained models.
3. Designing a two-stage training strategy to adapt to limited and noisy labeled samples.
4. Demonstrating the effectiveness of LLMTTT on various OOD graphs through experiments and theoretical analysis.

The researchers showcase the significant performance improvement of LLMTTT compared to existing OOD generalization methods through experiments and provide theoretical analysis to ensure its effectiveness.