GraphPrompt: Unifying Pre-Training and Downstream Tasks
for Graph Neural Networks

In particular, the downstream tasks are given limited supervision in a few-shot setting: for each class in the two tasks, only $k$ labeled samples (i.e., nodes or graphs) are provided, known as $k$-shot classification.

The success of GNNs boils down to the message-passing mechanism (Wu et al., 2020), in which each node receives and aggregates messages (i.e., features or embeddings) from its neighboring nodes to generate its own representation. This operation of neighborhood aggregation can be stacked in multiple layers to enable recursive message passing. Formally, in the $l$-th GNN layer, the embedding of node $v$, denoted by $\mathbf{h}^{l}_{v}$, is calculated based on the embeddings in the previous layer, as follows.

where $\mathcal{N}_{v}$ is the set of neighboring nodes of $v$, $\theta^{l}$ is the learnable GNN parameters in layer $l$. $\textsc{Aggr}(\cdot)$ is the neighborhood aggregation function and can take various forms, ranging from the simple mean pooling (Kipf and Welling, 2017; Hamilton et al., 2017) to advanced neural networks such as neural attention (Veličković et al., 2018) or multi-layer perceptrons (Xu et al., 2019). Note that in the first layer, the input node embedding $\mathbf{h}^{0}_{v}$ can be initialized as the node features in $\mathbf{X}$. The total learnable GNN parameters can be denoted as $\Theta=\{\theta^{1},\theta^{2},\ldots\}$. For brevity, we simply denote the output node representations of the last layer as $\mathbf{h}_{v}$.

We first introduce the overall framework of GraphPrompt in Fig. 2. Our framework is deployed on a set of label-free graphs shown in Fig. 2(a), for pre-training in Fig. 2(b). The pre-training adopts a link prediction task, which is self-supervised without requiring extra annotation. Afterward, in Fig. 2(c), we capitalize on a learnable prompt to guide each downstream task, namely, node classification or graph classification, for task-specific exploitation of the pre-trained model. In the following, we explain how the framework supports a unified view of pre-training and downstream tasks.

In comparison to other fields such as visual and language processing, graph learning is uniquely characterized by the exploitation of graph topology. In particular, subgraph is a universal structure capable of expressing both node- and graph-level instances. On one hand, at the node level, every node resides in a local neighborhood, which in turn contextualizes the node (Liu et al., 2020, 2021a, 2021c). The local neighborhood of a node $v$ on a graph $G=(V,E)$ is usually defined by a contextual subgraph $S_{v}=(V(S_{v}),E(S_{v}))$, where its set of nodes and edges are respectively given by

where $d(u,v)$ gives the shortest distance between nodes $u$ and $v$ on the graph $G$, and $\delta$ is a predetermined threshold. That is, $S_{v}$ consists of nodes within $\delta$ hops from the node $v$, and the edges between those nodes. Thus, the contextual subgraph $S_{v}$ embodies not only the self-information of the node $v$, but also rich contextual information to complement the self-information (Zhang and Chen, 2018; Huang and Zitnik, 2020). On the other hand, at the graph level, the maximum subgraph of a graph $G$, denoted $S_{G}$, is the graph itself, i.e., $S_{G}=G$. The maximum subgraph $S_{G}$ spontaneously embodies all information of $G$. In summary, subgraphs can be used to represent both node- and graph-level instances: Given an instance $x$ which can either be a node or a graph (e.g., $x=v$ or $x=G$), the subgraph $S_{x}$ offers a unified access to the information associated with $x$.

It is worth noting that node and graph classification can be further condensed into a single set of notations. Let $(x,y)$ be an annotated instance of graph data, i.e., $x$ is either a node or a graph, and $y\in Y$ is the class label of $x$ among the set of classes $Y$. Then,

Finally, to materialize the common task template, we discuss how to learn the subgraph embedding vector $\mathbf{s}_{x}$ for the subgraph $S_{x}$. Given node representations $\mathbf{h}_{v}$ generated by a GNN (see Sect. 3.2), a standard approach of computing $\mathbf{s}_{x}$ is to employ a ReadOut operation that aggregates the representations of nodes in the subgraph $S_{x}$. That is,

The choice of the aggregation scheme for ReadOut is flexible, including sum pooling and more advanced techniques (Ying et al., 2018; Xu et al., 2019). In our implementation, we simply use sum pooling.

In summary, the unification framework is enabled by the common task template of subgraph similarity learning, which lays the foundation of our pre-training and prompting strategies as we will introduce in the following parts.

As discussed earlier, our pre-training phase employs the link prediction task. Using link prediction/generation is a popular and natural way (Hamilton et al., 2017; Hu et al., 2020a; Hwang et al., 2020; Lu et al., 2021), as a vast number of links are readily available on large-scale graph data without extra annotation. In other words, the link prediction objective can be optimized on label-free graphs, such as those shown in Fig. 2(a), in a self-supervised manner.

Based on the common template defined in Sect. 4.1, the link prediction task is anchored on the similarity of the contextual subgraphs of two candidate nodes. Generally, the subgraphs of two positive (i.e., linked) candidates shall be more similar than those of negative (i.e., non-linked) candidates, as illustrated in Fig. 2(b). Subsequently, the pre-trained prior on subgraph similarity can be naturally transferred to node classification downstream, which shares a similar intuition: the subgraphs of nodes in the same class shall be more similar than those of nodes from different classes. On the other hand, the prior can also support graph classification downstream, as graph similarity is consistent with subgraph similarity not only in letter (as a graph is technically always a subgraph of itself), but also in spirit. The “spirit” here refers to the tendency that graphs sharing similar subgraphs are likely to be similar themselves, which means graph similarity can be translated into the similarity of the containing subgraphs (Shervashidze et al., 2011; Zhang et al., 2018; Togninalli et al., 2019).

Formally, given a node $v$ on graph $G$, we randomly sample one positive node $a$ from $v$’s neighbors, and a negative node $b$ from the graph that does not link to $v$, forming a triplet $(v,a,b)$. Our objective is to increase the similarity between the contextual subgraphs $S_{v}$ and $S_{a}$, while decreasing that between $S_{v}$ and $S_{b}$. More generally, on a set of label-free graphs $\mathcal{G}$, we sample a number of triplets from each graph to construct an overall training set $\mathcal{T}_{\text{pre}}$. Then, we define the following pre-training loss.

where $\tau$ is a temperature hyperparameter to control the shape of the output distribution. Note that the loss is parameterized by $\Theta$, which represents the GNN model weights.

The output of the pre-training phase is the optimal model parameters $\Theta_{0}=\arg\min_{\Theta}\mathcal{L}_{\text{pre}}(\Theta)$. $\Theta_{0}$ can be used to initialize the GNN weights for downstream tasks, thus enabling the transfer of prior knowledge downstream.

The unification of pre-training and downstream tasks enables more effective knowledge transfer as the tasks in the two phases are made more compatible by following a common template. However, it is still important to distinguish different downstream tasks, in order to capture task individuality and achieve task-specific optimum.

To cope with this challenge, we propose a novel task-specific learnable prompt on graphs, inspired by prompting in natural language processing (Brown et al., 2020). In language contexts, a prompt is initially a handcrafted instruction to guide the downstream task, which provides task-specific cues to extract relevant prior knowledge through a unified task template (typically, pre-training and downstream tasks are all mapped to masked language modeling). More recently, learnable prompts (Liu et al., 2021e; Lester et al., 2021) have been proposed as an alternative to handcrafted prompts, to alleviate the high engineering cost of the latter.

Formally, let $\mathbf{p}_{t}$ denote a learnable prompt vector for a downstream task $t$, as shown in Fig. 2(c). The prompt-assisted ReadOut operation on a subgraph $S_{x}$ for task $t$ is

where $\mathbf{s}_{t,x}$ is the task $t$-specific subgraph representation, and $\odot$ denotes the element-wise multiplication. That is, we perform a feature weighted summation of the node representations from the subgraph, where the prompt vector $\mathbf{p}_{t}$ is a dimension-wise reweighting in order to extract the most relevant prior knowledge for the task $t$.

Note that other prompt designs are also possible. For example, we could consider a learnable prompt matrix $\mathbf{P}_{t}$, which applies a linear transformation to the node representations:

More complex prompts such as an attention layer is another alternative. However, one of the main motivation of prompting instead of fine-tuning is to reduce reliance on labeled data. In few-shot settings, given very limited supervision, prompts with fewer parameters are preferred to mitigate the risk of overfitting. Hence, the feature weighting scheme in Eq. (12) is adopted for our prompting as the prompt is a single vector of the same length as the node representation, which is typically a small number (e.g., 128).

Formally, consider a task $t$ with a labeled training set $\mathcal{T}_{t}=\{(x_{1},y_{1}),(x_{2},y_{2}),\ldots\}$, where $x_{i}$ is an instance (i.e., a node or a graph), and $y_{i}\in Y$ is the class label of $x_{i}$ among the set of classes $Y$. The loss for prompt tuning is defined as

where the class prototypical subgraph for class $c$ is represented by $\tilde{\mathbf{s}}_{t,c}$, which is also generated by the prompt-assisted, task-specific ReadOut.

Note that, the prompt tuning loss is only parameterized by the learnable prompt vector $\mathbf{p}_{t}$, without the GNN weights. Instead, the pre-trained GNN weights $\Theta_{0}$ are frozen for downstream tasks, as no fine-tuning is necessary. This significantly decreases the number of parameters to be updated downstream, thus not only improving the computational efficiency of task learning and inference, but also reducing the reliance on labeled data.

We summarize these datasets in Table 1, and present further details in Appendix B. Note that the “Task” column indicates the type of downstream task performed on each dataset: “N” for node classification and “G” for graph classification.

Note that other few-shot learning methods on graphs, such as Meta-GNN (Zhou et al., 2019) and RALE (Liu et al., 2021b), adopt a meta-learning paradigm (Finn et al., 2017). Thus, they cannot be used in our setting, as they require labeled data in their base classes for the meta-training phase. In our approach, only label-free graphs are utilized for pre-training.

The downstream tasks follow a $k$-shot classification setting. For each type of downstream task, we construct a series of $k$-shot classification tasks. The details of task construction will be elaborated later when reporting the results in Sect. 5.2. For task evaluation, as the $k$-shot tasks are balanced classification, we employ accuracy as the evaluation metric following earlier work (Wang et al., 2020; Liu et al., 2021b).

For all the baselines, based on the authors’ code and default settings, we further tune their hyper-parameters to optimize their performance. We present more implementation details of the baselines and our GraphPrompt in Appendix D.

As discussed, we perform two types of downstream task different from the link prediction task in pre-training, namely, node classification and graph classification in few-shot settings. We first evaluate on a fixed-shot setting, and then vary the shot numbers to see the performance trend.

Table 2 illustrates the results of few-shot node classification. We have the following observations. First, our proposed GraphPrompt outperforms all the baselines across the three datasets, demonstrating the effectiveness of GraphPrompt in transferring knowledge from the pre-training to downstream tasks. In particular, by virtue of the unification framework and prompt-based task-specific aggregation in ReadOut function, GraphPrompt is able to narrow the gap between pre-training and downstream tasks, and guide the downstream tasks to exploit the pre-trained model in a task-specific manner. Second, compared to graph pre-training models, end-to-end GNN models can sometimes achieve comparable or even better performance. This implies that the discrepancy between the pre-training and downstream tasks in these pre-training approaches obstructs the knowledge transfer from the former to the latter. In such a case, even with sophisticated pre-training, they cannot effectively promote the performance of downstream tasks. Third, the graph prompt model GPPT is only comparable to or even worse than the other baselines, despite also using prompts. A potential reason is that GPPT requires much more learnable parameters in their prompts than ours, which may not work well given very few shots (e.g., 1-shot).

We illustrate the results of few-shot graph classification in Table 3, and have the following observations. First, our proposed GraphPrompt significantly outperforms the baselines on these four datasets. This again demonstrates the necessity of unification for pre-training and downstream tasks, and the effectiveness of prompt-assisted task-specific aggregation for ReadOut. Second, as both node and graph classification tasks share the same pre-trained model on PROTEINS and ENZYMES, the superior performance of GraphPrompt on both types of task further demonstrates that, the gap between different tasks is well addressed by virtue of our unification framework. Third, the graph pre-training models generally achieve better performance than the end-to-end GNN models. This is because both InfoGraph and GraphCL capitalize on graph-level tasks for pre-training, which are naturally closer to the downstream graph classification.

In general, our proposed GraphPrompt consistently outperforms the baselines especially with lower shots. For node classification, as the number of nodes in each graph is relatively small, 10 shots per class might be sufficient for semi-supervised node classification. Nevertheless, GraphPrompt is competitive even with 10 shots. For graph classification, GraphPrompt can be surpassed by some baselines when given more shots (e.g., 20 or more), especially on ENZYMES. On this dataset, 30 shots per class implies 30% of the 600 graphs are used for training, which is not our target scenario.

We further analyse several aspects of our model. Due to space constraint, we only report the ablation and parameter efficiency study, and leave the rest to Appendix E.

We conduct the ablation study on three datasets for node classification (Flickr, PROTEINS, and ENZYMES) and graph classification (COX2, ENZYMES, and BZR), respectively, and illustrate the comparison in Fig. 5. We have the following observations. (1) Without the prompt vector, no prompt usually performs the worst among the variants, showing the necessity of prompting the ReadOut operation differently for different downstream tasks. (2) Converting the prompt vector into a linear transformation matrix also hurts the performance, as the matrix involves more parameters thus increasing the reliance on labeled data.

In particular, as GIN works in an end-to-end manner, it is obvious that it involves the largest number of parameters for updating. For GPPT, it requires a separate learnable vector for each class as its representation, and an attention module to weigh the neighbors for aggregation in the structure token generation. Therefore, GPPT needs to update more parameters than GraphPrompt, which is one factor that impairs its performance in downstream tasks. For our proposed GraphPrompt, it not only outperforms the baselines GIN and GPPT as we have seen earlier, but also requires the least parameters and FLOPs for downstream tasks. For illustration, in addition to prompt tuning, if we also fine-tune the pre-trained weights instead of freezing them (denoted GraphPrompt+ft), there will be significantly more parameters to update.