Source-linked AI summary
Universal Prompt Tuning for Graph Neural Networks
Taoran Fang, Yunchao Zhang, Yang Yang, Chunping Wang, Lei Chen
TL;DR
Pre-trained GNN adaptation is complicated by diverse pre-training strategies and weaknesses of fine-tuning, including objective misalignment and few-shot overfitting. The paper proposes GPF, a feature-space prompt that works across strategies, and reports better performance than fine-tuning across full-shot and few-shot settings.
Problem
Graph prompt tuning lacks a broadly applicable method because graph pre-training strategies are diverse, while fine-tuning can suffer objective misalignment, catastrophic forgetting, and few-shot overfitting.
Method
GPF adds a shared learnable vector to input node features before a frozen pre-trained GNN, with GPF-plus providing node-specific vectors.
Results
Across pre-training strategies, GPF and GPF-plus outperform fine-tuning on average, with improvements of 1.14% and 1.60%, respectively.
Takeaways & Limitations
Feature-space graph prompts provide a universal alternative to strategy-specific prompting functions for downstream adaptation.
Takeaways & Limitations
The theoretical comparison with fine-tuning assumes a squared regression loss and includes a stated degeneracy condition.
Abstract
from arXiv · showhide
In recent years, prompt tuning has sparked a research surge in adapting pre-trained models. Unlike the unified pre-training strategy employed in the language field, the graph field exhibits diverse pre-training strategies, posing challenges in designing appropriate prompt-based tuning methods for graph neural networks. While some pioneering work has devised specialized prompting functions for models that employ edge prediction as their pre-training tasks, these methods are limited to specific pre-trained GNN models and lack broader applicability. In this paper, we introduce a universal prompt-based tuning method called Graph Prompt Feature (GPF) for pre-trained GNN models under any pre-training strategy. GPF operates on the input graph's feature space and can theoretically achieve an equivalent effect to any form of prompting function. Consequently, we no longer need to illustrate the prompting function corresponding to each pre-training strategy explicitly. Instead, we employ GPF to obtain the prompted graph for the downstream task in an adaptive manner. We provide rigorous derivations to demonstrate the universality of GPF and make guarantee of its effectiveness. The experimental results under various pre-training strategies indicate that our method performs better than fine-tuning, with an average improvement of about 1.4% in full-shot scenarios and about 3.2% in few-shot scenarios. Moreover, our method significantly outperforms existing specialized prompt-based tuning methods when applied to models utilizing the pre-training strategy they specialize in. These numerous advantages position our method as a compelling alternative to fine-tuning for downstream adaptations.
1 Introduction
Pre-trained GNNs face objective misalignment, catastrophic forgetting, and overfitting during downstream adaptation, while graph prompt tuning is difficult because pre-training strategies vary. GPF addresses this by tuning input features universally across pre-trained GNN models.
- 1 Introduction: Pre-train, fine-tune adaptation suffers from pre-training/downstream objective misalignment and can overfit small downstream datasets through catastrophic forgetting.These issues can render pre-training ineffective in few-shot settings.
- 1 Introduction: Prompt tuning freezes the pre-trained model and modifies input data instead of updating model parameters.This provides an alternative to fine-tuning for downstream adaptation.
- 1 Introduction: Graph prompting is challenging because prompting functions should align with diverse graph pre-training strategies, unlike the more unified language-model setting.Existing graph methods are therefore difficult to generalize across strategies.
- 1 Introduction: GPF operates in the input graph’s feature space and can match any prompting function across pre-trained GNN models.It adds a shared learnable vector to node features, avoiding strategy-specific prompting functions.
- 1 Introduction: 1.4% average improvement in full-shot scenarios and 3.2% in few-shot scenarios over fine-tuning are reported for GPF and GPF-plus.The methods also outperform specialized prompt-based methods on models using their target pre-training strategy.
2 Related work
Prior graph prompt-tuning methods adapt models with manually designed templates, but the field’s exploration remains limited and concentrated on edge-prediction pre-training.
- 2 Related work: Prompt-based tuning generally freezes pre-trained model parameters and adds learnable components in the input space.This input-space design has been used in language and vision prompting.
- 2 Related work: Existing graph prompt methods introduce virtual class-prototype nodes or graphs with learnable links, primarily for edge-prediction-pre-trained models.These templates make downstream adaptation resemble edge prediction.
- 2 Related work: Graph prompt tuning remains limited compared with prompting research in natural language processing and computer vision.The cited graph approaches focus on specialized settings rather than broad pre-training strategies.
3 Methodology
The methodology frames graph prompting as transforming downstream inputs before a frozen GNN, then proposes feature-space prompts designed to work across pre-training strategies. GPF adds one shared vector per graph, while GPF-plus assigns node-specific vectors, with theory supporting universality and competitive tuning.
- 3.2 Graph Prompt Tuning: The graph prompt pipeline transforms each input graph into a prompted graph, which is then processed by a frozen pre-trained GNN and projection head.Prompt optimization searches over prompted graph components to maximize downstream label likelihood.
- 3.2 Graph Prompt Tuning: Graph template design depends on the pre-training task, but existing prompting functions mainly address edge prediction and lack coverage for strategies such as attribute masking and context prediction.The paper motivates a universal feature-space alternative to strategy-specific template design.
- 3.3 Universal Graph Prompt Design: GPF adds a shared learnable feature vector to every node, replacing the original features before processing by the pre-trained GNN.The vector has the same dimension as node features.
- 3.3 Universal Graph Prompt Design: GPF-plus assigns an independent learnable vector to each node, increasing flexibility but requiring O(N) parameters and creating difficulties for varying graph sizes.The method uses a more parameter-intensive node-specific design than GPF.
- 3.4 Theoretical Analysis: GPF and GPF-plus are theoretically universal: they can achieve effects equivalent to arbitrary prompting functions for pre-trained GNNs.Theorem 1 states that an extra GPF vector exists for any prompted graph in the prompting template’s candidate space.
- 3.4 Theoretical Analysis: Under the stated theoretical analysis, GPF can obtain a lower minimum loss than fine-tuning in certain scenarios.The conclusion is tied to the assumptions used for the squared regression-loss analysis.
4 Experiments
Experiments evaluate fine-tuning and graph prompt methods across five pre-training strategies and chemistry and biology benchmarks. GPF and GPF-plus generally outperform fine-tuning and specialized graph prompting while using fewer tunable parameters.
- Experimental Setup: The study evaluates five pre-training strategies, including Infomax, Edge Prediction, Attribute Masking, Context Prediction, and Graph Contrastive Learning.Experiments use a 5-layer GIN and chemistry and biology datasets, with five random-seed runs per setting.
- 4.2 Main Results: GPF outperforms fine-tuning in 28/36 experiments, while GPF-plus outperforms fine-tuning in 29/36 experiments.The prompt methods use significantly fewer tunable parameters than fine-tuning.
- 4.2 Main Results: GPF achieves an average improvement of 1.14%, while GPF-plus achieves an average improvement of 1.60% over fine-tuning across the examined pre-training strategies.The experiments use five pre-training strategies and report average results across the evaluated settings.
- 4.2 Main Results: GPF-plus performs better than GPF in 26/36 experiments, consistent with its greater flexibility and expressiveness.GPF-plus uses learnable basis vectors and linear projections, whereas GPF uses an extra learnable feature vector.
- 4.3 Comparison with Existing Graph Prompt-based Methods: On chemistry and biology benchmarks, GPF and GPF-plus achieve average improvements of 12%, 3%, and 13% over GPPT, GPPT (w/o ol), and GraphPrompt, respectively.These comparisons use models pre-trained by Edge Prediction.
5 Conclusion
The paper concludes that GPF and GPF-plus provide universal prompt-based tuning for pre-trained GNNs by operating in downstream graph feature space. They theoretically match arbitrary prompting functions and empirically offer an alternative to fine-tuning, including for node-wise tasks.
- 5 Conclusion: GPF and GPF-plus operate on downstream input-graph features and can theoretically achieve an equivalent effect to any prompting function.This removes the need to specify a separate prompting function for each pre-training strategy.
- 5 Conclusion: The methods adapt pre-trained GNNs by obtaining a prompted graph for downstream tasks without explicitly matching each pre-training strategy.The conclusion presents this adaptive feature-space procedure as the central universal-tuning approach.
- A.1 Extension to node-wise tasks: For node classification and link prediction, Subgraph GNNs provide graph-level and node-level representations while the pre-trained model is frozen.A learnable graph prompt is introduced for induced subgraphs before downstream node-level tasks.
A.2 Proof for Theorem 1
Theorem 1 establishes that GPF can reproduce the effect of arbitrary prompting functions for pre-trained GNNs by modifying node features, including feature, link, and isolated-component transformations.
- GPF can achieve the same graph representation as any prompted graph generated by an arbitrary prompting function.The theorem imposes no constraint on the prompting function or the resulting adjacency and feature matrices.
- GPF covers arbitrary graph-level transformations by decomposing them into feature, link, and isolated-component transformations.These transformations respectively modify node features, adjacency structure, or graph scale through isolated subgraphs.
- GPF can reproduce arbitrary feature transformations and obtain the exact graph representation produced by the corresponding transformed input.The derivation compares the representation change from a general feature perturbation with that induced by the shared GPF vector.
- GPF adds one shared learnable feature vector to every node, yet can match independent node-wise feature modifications under the analyzed GNN architecture.The proof uses sum or other weighted aggregation readouts and extends from the single-layer linear case to multi-layer models.
- GPF also encompasses arbitrary link transformations, showing that feature-space and structural-space modifications can have equivalent effects in the analyzed GNNs.The result applies to adjacency changes represented by adding or removing edges.
- The universality result extends to isolated-component transformations and series of feature, link, and isolated-component transformations.Thus, GPF can replicate graph-level transformations that alter graph scale as well as local features or links.
A.3 Proof for Theorem 2
Theorem 2 shows that, under a non-degeneracy condition, GPF with a linear projection head can achieve lower theoretical tuning loss than fine-tuning.
- Under the stated non-degeneracy condition, the optimal GPF loss is strictly lower than the optimal fine-tuning loss.The proof states that l_GPF can attain its constructed solution while l_FT remains greater than 0, yielding l_GPF < l_FT.
- The analysis represents each graph through a diffusion matrix, node features modified by a shared vector, a frozen projection, and a linear prediction head.For graph G_i, the GPF representation is computed as S_i · (X_i + [1]_T · p) · W.
- The construction permits fine-tuning’s effective parameter change to be represented by an arbitrary vector while retaining a positive separation from the GPF solution.This separation is formalized through δ > 0 in the proposition used to prove the theorem.
- The proof compares GPF and fine-tuning by expressing graph representations over the set of unique node features and their graph-specific coefficients.The coefficient matrix captures how each graph aggregates the unique features in the dataset.
- The non-degeneracy assumption requires that the coefficient matrix be column full-rank, excluding a uniform feature distribution shared across graphs.This assumption is described as aligning with real-world scenarios in which graphs do not share an identical feature distribution.
B.1 Details of the datasets
The experiments use chemistry and biology pre-training data, downstream graph datasets, and five common self-supervised or contrastive pre-training strategies.
- The pre-training data cover chemistry and biology: 2 million unlabeled ZINC15 molecules plus 256K labeled ChEMBL molecules, and 395K unlabeled plus 88K labeled protein ego-networks.The biology data come from PPI networks, while the chemistry data combine molecular pre-training and property-prediction examples.
- The downstream datasets for models pre-trained on Biology and Chemistry are summarized in Table 3.The supplied passage identifies the table as reporting downstream-dataset statistics but does not provide its cell values.
- The study evaluates Deep Graph Infomax, Edge Prediction, Attribute Masking, Context Prediction, and Graph Contrastive Learning as pre-training strategies.These strategies span mutual-information maximization, graph reconstruction, masking, contextual prediction, and contrastive learning.
- For the first four strategies, pre-training follows Hu et al.’s procedure, while Graph Contrastive Learning follows You et al.’s procedure.The models also receive supervised graph-level property prediction to further enhance pre-trained-model performance.
B.3 Results of few-shot graph classification
In few-shot graph classification, graph prompt tuning outperforms fine-tuning across the evaluated pre-training strategies, with the strongest reported pattern in 100-shot experiments.
- Table 4 reports 50-shot test ROC-AUC (%) for molecular prediction and protein function prediction benchmarks.The supplied table passage identifies the scope and metric but does not provide individual benchmark values.
- 42 of 45 evaluated 100-shot cases are won by GPF or GPF-plus, and both methods exceed fine-tuning on average across all pre-training strategies.GPF is best in 14 cases and GPF-plus in 28 cases.
B.4 Parameter efficiency analysis
GPF and GPF-plus achieve substantial parameter efficiency compared with fine-tuning, while graph prompts improve adaptation performance in linear probing.
- B.4 Parameter efficiency analysis: GPF uses no more than 0.02% of fine-tuning’s tunable parameters, while GPF-plus uses no more than 0.7%.These comparisons exclude the task-specific projection head θ.
- B.4 Parameter efficiency analysis: GPF-based methods reduce the training time and storage space required for downstream adaptations.
- B.4 Parameter efficiency analysis: Integrating a graph prompt is important for effectively adapting pre-trained models compared with linear probing.
B.6 Comparison with other tuning methods
The proposed graph prompt tuning methods outperform the alternative tuning strategies evaluated on protein function prediction benchmarks.
- B.6 Comparison with other tuning methods: PARTIAL-k tunes the last k layers while freezing the remaining pre-trained model parameters.
- B.6 Comparison with other tuning methods: MLP-k freezes the pre-trained model and uses a k-layer MLP as the classification projection head.
- B.6 Comparison with other tuning methods: Our methods outperform other tuning methods in all cases on the biology datasets.The comparison includes the alternative tuning strategies described for partial-layer tuning and MLP projection heads.
B.7 Extra results on GCC
On graph classification benchmarks using GCC-pre-trained models, GPF consistently outperforms fine-tuning on IMDB-BINARY and IMDB-MULTI.
- B.7 Extra results on GCC: GCC pre-training uses a self-supervised graph contrastive learning strategy on six graph datasets.
- B.7 Extra results on GCC: GPF consistently outperforms fine-tuning when adapting GCC-pre-trained models on IMDB-BINARY and IMDB-MULTI.The results are reported for graph classification benchmarks.
- B.7 Extra results on GCC: The GCC downstream evaluation uses IMDB-BINARY and IMDB-MULTI, whose graphs have associated target labels.