All in One: Multi-task Prompting for Graph Neural Networks (Extended Abstract)

1 Introduction

Graph neural networks (GNNs) [Sun et al., 2021a] are increasingly applied across various fields [Sun et al., 2023b; Sun et al., 2022b; Li et al., 2024a; Sun et al., 2022c; Chen et al., 2020; Sun et al., 2023a]. The focus has shifted towards optimizing graph model training for specific problems. Traditional graph learning methods depend heavily on labels, often scarce or unfit for real-world complexities, leading to overfitting, especially with out-of-distribution data [Shen et al., 2021]. A popular mitigation strategy involves pre-training on accessible data, then fine-tuning for specific tasks [Jin et al., 2020], despite challenges in aligning pre-trained models with diverse downstream tasks.

A novel approach, inspired by NLP, combines pre-training with prompt learning and fine-tuning, where prompts facilitate task-specific model adjustments without extensive retraining. This method shows promise for efficient model adaptation, especially in scenarios with limited data. However, applying the concept of language prompts to GNNs introduces challenges, such as defining prompt content and integration with graph structures, and ensuring prompts effectively bridge pre-training tasks with varied downstream applications. Current efforts in graph prompt learning are limited and typically focus on single-task scenarios [Sun et al., 2022a]. We extend NLP prompt methods to GNNs for multitask applications, addressing challenges in prompt design, task reformulation, and prompt optimization. Our contributions include a unified prompt format for language and graph domains, a strategy to reformulate tasks for better alignment with pre-training, and the application of meta-learning to enhance prompt efficacy across multiple tasks. Our extensive evaluations demonstrate the superiority of our approach.

∗The original paper was published in KDD23

2 Motivations

Graph pre-training [Sun et al., 2021b] employs strategies to imbue GNNs with broad knowledge, reducing the need for task-specific annotations. Techniques vary from node and edge comparisons to graph-level contrastive learning, which proves superior in learning graph knowledge by enhancing graph representation or adjusting model parameters for consistency across perturbations [You et al., 2020; Xia et al., 2022; Sun et al., 2023b]. Intuitively, the above graph-level pre-training strategies have some intrinsic similarities with the language-masked prediction task: aligning two graph views generated by node/edge/feature mask or other perturbations is very similar to predicting some vacant "blanks" on graphs. To this end, we aim to merge graph pretraining's depth with prompt learning's adaptability, addressing the multifaceted challenges in deploying GNNs across various tasks more effectively.

3 Multi-task Prompting on Graphs

This section presents a condensed overview of our approach to multi-task prompting for graph models, aiming to enhance the transferability of pre-trained graph models across various tasks without altering the original model architecture.

Objective: Our primary goal is to develop a graph prompt that seamlessly integrates with original graphs, thereby aligning pre-trained graph models more closely with diverse downstream tasks and improving knowledge transfer across domains.

FrameworkOverview: We introduce a multi-task prompting framework that first standardizes different graph tasks into a uniform format, focusing on graph-level tasks. We then design a novel graph prompt that incorporates learnable tokens, structures, and adaptive insertion patterns. To optimize the prompt for various tasks, we employ a meta-learning strategy, enabling the framework to adjust prompts dynamically for improved performance across multiple tasks.

Reformulating Tasks for Generalization: Recognizing the challenge of diverse task requirements in graphs, we reformulate node-level and edge-level tasks into graph-level tasks. This approach, inspired by the hierarchical nature of graph operations, allows for a broader application of pretraining knowledge by treating operations like node or edge modifications as graph-level changes.

Designing the Prompt Graph: We draw parallels between NLP and graph prompting, aiming for a unified representation that includes prompt tokens, token structures, and insertion patterns. This ensures that our graph prompts are both meaningful and adaptable to the structure of the original graph.

Let a graph instance be G = (V, E) where V = {v1, v2, · · · , vN} is the node set containing N nodes; each node has a feature vector denoted by xi ∈ R1×d for node vi; E = {(vi, vj)|vi, vj ∈ V} is the edge set where each edge connects a pair of nodes in V. With the previous discussion, we here present our prompt graph as Gp = (P, S) where P = {p1, p2, · · · , p|P|} denotes the set of prompt tokens and |P| is the number of tokens. Each token pi ∈ P can be represented by a token vector pi ∈ R1×d with the same size of node features in the input graph; Note that in practice, we usually have |P| ≪ N and |P| ≪ dh where dh is the size of the hidden layer in the pre-trained graph model. With these token vectors, the input graph can be reformulated by adding the j-th token to graph node vi (e.g., ˆxi = xi + pj). Then, we replace the input features with the prompted features and send them to the pre-trained model for further processing.

S = {(pi, pj)|pi, pj ∈ P} is the token structure denoted by pair-wise relations among tokens. Unlike the NLP prompt, the token structure in the prompt graph is usually implicit. To solve this problem, we propose three methods to design the prompt token structures: (1) the first way is to learn tunable parameters:

where aij is a tunable parameter indicating how possible the token pi and the token pj should be connected; (2) the second way is to use the dot product of each prompt token pair and prune them according to the dot value. In this case, (pi, pj) ∈ S iff σ(pi · pj) < δ where σ(·) is a sigmoid function and δ is a pre-defined threshold; (3) the third way is to treat the tokens as independent and then we have S = ∅.

Let ψ be the inserting function that indicates how to add the prompt graph Gp to the input graph G, then the manipulated graph can be denoted as Gm = ψ(G, Gp). We can define the inserting pattern as the dot product between prompt tokens and input graph nodes, and then use a tailored connection like ˆxi = xi + ∑|P| k=1 wikpk where wik is a weighted value to prune unnecessary connections:

As an alternative and special case, we can also use a more simplified way to get ˆxi = xi + ∑|P| k=1 pk.

Meta-Learning for Prompt Optimization: We leverage meta-learning to refine our prompting approach, structuring the learning process to accommodate multiple tasks simultaneously. This method updates prompt parameters based on task-specific performances, ensuring that the final prompts are well-suited to a wide array of graph tasks.

(b) Induced graphs for edges

Figure 2: Reformulate node-level and edge-level tasks to graph-level tasks by induced graphs.

4 Why It Works?

Comparison to Prior Work: While GPPT [Sun et al., 2022a] represents an early attempt at graph prompting, focusing on edge prediction for node classification, our method extends this concept significantly. Unlike GPPT, our framework is more versatile, accommodating a broader range of graph tasks and pre-training strategies beyond edge prediction, including advanced graph-level strategies like GraphCL [You et al., 2020] and SimGRACE [Xia et al., 2022].

Flexibility: Our approach introduces the concept of a prompt graph comprising multiple tokens with learnable structures, offering a more nuanced and flexible method for graph manipulation to better align with various pre-training strategies. We demonstrate that this flexibility allows for more effective adaptations of the graph structure to suit different tasks, reducing the error margin in representing manipulated graphs.

The nature of prompting is to manipulate the input data to match the pretext. Therefore, the flexibility of data operations is the bottleneck of prompting performance. Let g be any graph-level transformation such as "changing node features", "adding or removing edges/subgraphs" etc., and φ∗ be the frozen pre-trained graph model. For any graph G with adjacency matrix A and node feature matrix X, Fang et al. [Fang et al., 2022] have proved that we can always learn an appropriate prompt token p∗ making the following equation stand:

This means we can learn an appropriate token applied to the original graph to imitate any graph manipulation. Here Opφ denotes the error bound between the manipulated graph and the prompting graph w.r.t. their representations from the pretrained graph model. This error bound is related to some nonlinear layers of the model (unchangeable) and the quality of the learned prompt (changeable), which is promising to be further narrowed down by a more advanced prompt scheme. In this paper, we extend the standalone token to a prompt graph that has multiple prompt tokens with learnable inner structures. Unlike the indiscriminate inserting in Equation (2) ("X+ p∗ " means the prompt token should be added to every node of the original graph), the inserting pattern of our proposed prompt graph is highly customized. Let ψ(G, Gp) denote the inserting pattern defined in section 3; G is the original graph, and Gp is the prompt graph, then we can learn an optimal prompt graph G ∗ p to extend Equation (2) as follows:

By efficient tuning, the new error bound O∗ pφ can be further reduced. That means our method supports more flexible transformations on graphs to match various pre-training strategies.

5 Evaluation

We compare our methods with other approaches on five public datasets including Cora [Welling and Kipf, 2016], CiteSeer [Welling and Kipf, 2016], Reddit [Hamilton et al., 2017], Amazon [Shchur et al., 2018], and Pubmed [Welling and Kipf, 2016]. We compare our method with supervised, pre-training plus fine-tuning, and other prompt methods across node, edge, and graph-level tasks. Key findings include our method's superior performance in multi-task settings, showcasing notable improvements over existing methods.

Multi-Task Performance Our study evaluates the performance of our prompt-based method across node-level, edgelevel, and graph-level tasks in few-shot learning settings, comparing it against supervised methods and pre-training approaches. Results in Table 1 show that supervised methods struggle due to limited annotations available in fewshot scenarios, while pre-training methods offer better performance by leveraging prior knowledge. However, selecting and fine-tuning a pre-trained model for a specific task is effort-intensive and not always transferable to other tasks.

Our method, by incorporating prompts, shows compatibility improvements across all task levels, achieving performance boosts ranging from 1.10% to 8.81% for node-level tasks, 1.28% to 12.26% for edge-level tasks, and 0.14% to 10.77% for graph-level tasks. Notably, our approach under a more challenging setting (with only 100 labeled samples per class) still outperforms the GPPT model, which uses a 30% to 50% label ratio, indicating superior efficiency and adaptability of our method in few-shot learning contexts across various graph tasks. Please see the original paper for more task performance like edge-level and graph-level tasks.

Transferability: Our method demonstrates enhanced adaptability, outperforming both hard transfer and finetuning approaches in transferring models to new tasks (as shown in Table 2) and domains (as shown in Table 3). This is particularly evident in tasks requiring significant adaptation, where our prompting framework facilitates more effective knowledge transfer.

Graph Transformation Flexibility: Our approach effectively minimizes the error in representing manipulated graphs, demonstrating its capacity to support a wide range of graph transformations. This is further illustrated by visualizations that highlight the improved graph representations achieved through our prompting method.

For more experiments, please see in our original paper.

6 A Bigger Picture of Graph Prompts

In the rapidly evolving field of Artificial General Intelligence (AGI) [Li et al., 2024b], significant advancements have been made, especially with applications like ChatGPT in NLP and Midjourney in Computer Vision (CV), greatly enhancing our efficiency and creativity. Yet, the application of AGI in graph data analysis remains nascent, despite its potential to revolutionize areas such as drug design and battery development due to challenges in harmonizing information across modalities, domains, and tasks.

In section 4, we can find that graph prompt has the potential to simulate various data manipulations. This means it can be used to achieve the tough challenge in graph domain transfer. Our paper also demonstrates its huge potential for task transfer. In addition, since we only need to tune a light-weight prompt while keeping a large graph model unchanged, it is more efficient. Prompt learning emerges as a promising solution. It has shown remarkable success in NLP and CV by reformulating tasks to leverage pre-trained models without extensive tuning. Prompt learning's efficiency in knowledge extraction and task reformulation presents an opportunity to address the complexities of working with graph data, suggesting a path to extend its benefits to graph-based AGI applications.

To achieve this end, some recent works have been proposed to follow up our paper [Sun et al., 2023c]. We recently further studied the feasibility of domain transferring with graph prompt [Zhao et al., 2024], explored the application of graph prompt in protein multimer structure prediction [Gao et al., 2024], and proposed various variants of graph prompt [Chen et al., 2024]. We also release "ProG" (Prompt Graph), which is a Python library built upon PyTorch to easily conduct single or multi-task prompting for pre-trained Graph Neural Networks (GNNs). Please use the library at https://github.com/sheldonresearch/ProG/

In future work, we can further study the integration of graph prompts with various graph models [Zhang et al., 2022b], extend its applications[Liu et al., 2022; Piao et al., 2023; Cui et al., 2023; Meng et al., 2023], and the security issue of graph prompt[Zhang et al., 2022a; Yang et al., 2023].

Acknowledgments

This research is supported in part by project #MMT-p2-23 of the Shun Hing Institute of Advanced Engineering, The Chinese University of Hong Kong, by grants from the Research Grant Council of the Hong Kong Special Administrative Region, China (No. CUHK 14217622), NSFC (No. 61972087, No. 62206067, No. U1936205, No. 62172300, No. 62202336), Guangzhou-HKUST(GZ) Joint Funding Scheme (No. 2023A03J0673), National Key R&D Program of China (No. 2022YFB3104300, No. 2021YFC3300300), the Fundamental Research Funds for the Central Universities (No. ZD-21-202101), and Open Research Projects of Zhejiang Lab (No. 2021KH0AB04).