RELIEF: Reinforcement Learning Empowered Graph Feature Prompt Tuning

TL;DR

RELIEF introduces a reinforcement learning-based method for strategically selecting and attaching lightweight feature prompts to specific graph nodes, significantly improving graph classification performance and data efficiency in few-shot scenarios.

Contribution

The paper proposes RELIEF, a novel RL-based approach for optimizing node prompt selection and content, enhancing graph prompt tuning beyond existing methods.

Findings

Outperforms fine-tuning and other prompt methods in classification accuracy.

Enhances data efficiency in few-shot learning scenarios.

Demonstrates effectiveness across various pre-training strategies.

Abstract

The advent of the "pre-train, prompt" paradigm has recently extended its generalization ability and data efficiency to graph representation learning, following its achievements in Natural Language Processing (NLP). Initial graph prompt tuning approaches tailored specialized prompting functions for Graph Neural Network (GNN) models pre-trained with specific strategies, such as edge prediction, thus limiting their applicability. In contrast, another pioneering line of research has explored universal prompting via adding prompts to the input graph's feature space, thereby removing the reliance on specific pre-training strategies. However, the necessity to add feature prompts to all nodes remains an open question. Motivated by findings from prompt tuning research in the NLP domain, which suggest that highly capable pre-trained models need less conditioning signal to achieve desired behaviors, we advocate for strategically incorporating necessary and lightweight feature prompts to certain graph nodes to enhance downstream task performance. This introduces a combinatorial optimization problem, requiring a policy to decide 1) which nodes to prompt and 2) what specific feature prompts to attach. We then address the problem by framing the prompt incorporation process as a sequential decision-making problem and propose our method, RELIEF, which employs Reinforcement Learning (RL) to optimize it. At each step, the RL agent selects a node (discrete action) and determines the prompt content (continuous action), aiming to maximize cumulative performance gain. Extensive experiments on graph and node-level tasks with various pre-training strategies in few-shot scenarios demonstrate that our RELIEF outperforms fine-tuning and other prompt-based approaches in classification performance and data efficiency. The code is available at https://github.com/JasonZhujp/RELIEF.