HGMP: Heterogeneous Graph Multi-Task Prompt Learning

1 Introduction

Heterogeneous graphs are powerful data structures that can simultaneously represent multiple types of nodes and edges, making them ideal for modeling complex relationships in various domains. For example, in social networks [Cao et al., 2021], heterogeneous graphs can model different types of entities such as users, posts, and comments, as well as various types of relationships, including friendships, comments, and likes. This flexibility enables heterogeneous graphs to capture rich and multifaceted interactions between entities, making them highly suitable for applications that require integration and analysis of diverse information, such as knowledge graphs [Ji et al., 2021; Hogan et al., 2021], recommender systems [Gao et al., 2023], and citation networks.

∗Corresponding author.

Graph Neural Networks (GNNs) [Yun et al., 2019; Kipf and Welling, 2016; Veliˇckovi´c et al., 2017] have demonstrated great performance in learning representations of graph-structured data by capturing local node relationships through message passing and aggregation. Heterogeneous Graph Neural Networks (HGNNs) [Mao et al., 2023; Fu et al., 2020; Yang et al., 2020] are an extension of GNNs, designed to handle the complexities of heterogeneous graphs. HGNNs incorporate both node and edge type information, allowing message passing to be more specialized and more effectively capture heterogeneous relationships. These models show improvements over standard GNNs in heterogeneous environments, particularly for tasks where understanding the relationships between entities depends heavily on node and edge types. However, the performance of both standard GNNs and HGNNs is limited by the scarcity of labeled data and the generalization ability on out-of-distribution data [Hu et al., 2019; Knyazev et al., 2019; Yehudai et al., 2021; Morris et al., 2019]. Furthermore, in multi-task scenarios, each task requires separate training, which further increases computational complexity.

To address these limitations, pre-training and fine-tuning strategies [Dong et al., 2019] have been widely adopted. During pre-training, models learn general feature representations from large-scale unlabeled data, reducing the dependence on labeled datasets and improving generalization. Finetuning is then applied to adapt the model to specific tasks. However, the mismatch between pre-training objectives and downstream task requirements can negatively impact performance [Liu et al., 2023], especially in scenarios with label imbalance or few-shot learning, where models struggle to generalize to new tasks or unseen data [Zhang et al., 2022].

Prompt learning, a rapidly emerging technique in the machine learning community, provides an attractive alternative. It embeds task-specific prompts into the input data, guiding the model to perform specific tasks. For example, a sentiment analysis task like "IJCAI2025 will bring together top researchers from around the world. I feel so [MASK]" can be reformulated using a preset prompt ("I feel so [MASK]"). It is anticipated that a language model trained on masked word prediction might predict "[MASK]" as "enthusiastic" rather than "disappointed," without requiring further parameter tuning for the new sentiment task, due to the model's pre-training that embeds relevant knowledge to address this type of question. This demonstrates that prompt learning can effectively leverage existing knowledge, enabling the model to generalize across tasks with minimal additional training.

The success of prompt learning in NLP [Li and Liang, 2021; Lester et al., 2021] has inspired exploration into its application in graph learning. However, in graph learning, the use of prompt learning is still in its early stages. While some early works have explored the use of prompts in homogeneous graph settings [Fang et al., 2024; Sun et al., 2023a; Sun et al., 2022; Liu et al., 2023; Tan et al., 2023; Huang et al., 2024; Sun et al., 2023b; Zi et al., 2024], these approaches do not address the unique challenges posed by heterogeneous graphs. Even recent prompt learning methods specifically designed for heterogeneous graphs still face some challenges [Yu et al., 2024; Ma et al., 2024]. First, aligning the pre-training objectives with the requirements of downstream tasks remains complex, as many early methods focus on link prediction tasks that connect data, neglecting the richness of heterogeneous data. Additionally, methods designed for single-task settings limit their applicability in multi-task scenarios. The main challenges in heterogeneous graph prompt learning can be summarized as: (C1) how to design unified graph templates for multi-task scenarios, (C2) how to develop pre-training methods that adapt to heterogeneous data, and (C3) how to design prompts that account for heterogeneity and feature diversity.

To address these challenges, we propose the HGMP (Heterogeneous Graph Multi-task Prompt Learning) framework. Our approach rephrases node and edge-level tasks as graphlevel tasks, aligning both downstream tasks and pre-training tasks in a common graph-level task space. This redefinition of tasks allows the model to undergo unified and effective training on graph-level objectives, capturing global graph structural features and enhancing the model's generalization ability. We introduce a heterogeneous pre-training strategy that combines subgraph extraction and augmentation for contrastive pre-training, fully leveraging the rich and diverse information in heterogeneous graphs. Furthermore, we design node-type-specific prompts using shared learnable vectors to modify node features, thereby enhancing the representation of input data and making it more suitable for the pre-training objectives.

In summary, the contributions of this paper are as follows:

• We propose a method that transforms node-level and edge-level tasks into graph-level tasks and introduces a flexible heterogeneous graph contrastive pre-training strategy to better capture the complexities of heterogeneous graph data.
• We design the HGMP framework, specifically tailored for multi-task heterogeneous graph scenarios, which effectively bridges the gap between pre-training and
• downstream task objectives, enhancing the model's generalization ability.
• We conduct comprehensive experiments on multiple public datasets, demonstrating the effectiveness, robustness, and adaptability of the HGMP framework across various task settings.

2 RelatedWork

Heterogeneous Graph Neural Networks. GNNs were initially designed for homogeneous graphs, focusing on node representation learning in graphs with consistent structures and attributes. To address the complexity of heterogeneous graphs, HGNNs emerged, enabling the modeling of diverse node and edge types. Representative works include HAN [Wang et al., 2019], which employs hierarchical attention mechanisms, HGT [Hu et al., 2020], a Transformerbased architecture for large-scale graphs, and SimpleHGN [Lv et al., 2021], which enhances GAT [Veliˇckovi´c et al., 2017] with learnable edge-type embeddings. Despite these advancements, HGNNs heavily rely on high-quality labeled data, which is often costly or infeasible to obtain, driving interest in pre-training techniques as alternatives to supervised learning.

Graph Pre-training. Graph pre-training aims to learn general knowledge from readily available data, reducing annotation costs. For homogeneous graphs, methods like GraphCL [You et al., 2020] utilize contrastive learning to exploit graph structures. In heterogeneous graphs, pre-training methods must handle diverse semantic relationships. Generative approaches reconstruct masked structures, while contrastive learning maximizes mutual information between positive sample pairs [Jing et al., 2021; Park et al., 2020], often outperforming generative methods. However, misalignment between pre-training objectives and downstream tasks can hinder the effective utilization of pre-trained knowledge. Prompt-based Learning On Graphs. Prompt learning bridges the gap between pre-training and downstream tasks by introducing task-specific prompts, reducing finetuning demands and enhancing adaptability. Methods like GPPT [Sun et al., 2022], GPF [Fang et al., 2024], and ProG [Sun et al., 2023a] leverage prompts for task alignment and adaptation. In heterogeneous graphs, approaches such as HGPrompt [Yu et al., 2024] and HetGPT [Ma et al., 2024] use templates to support few-shot learning and align pre-training with downstream tasks. Nonetheless, current graph prompt learning methods often struggle with multi-task demands or rely on rigid templates, limiting flexibility and effectiveness.

3 Preliminaries

Heterogeneous Graph. In graph data structures, a Heterogeneous Graph is an extended form designed to model complex networks with different types of nodes and edges. Compared to a homogeneous graph, a heterogeneous graph can more accurately represent diverse entities and their relationships in the real world. Formally, a heterogeneous graph can be defined as G = (V,E, TV , TE), where V represents the set of nodes, E represents the set of edges, and TV and TE denote the sets of node types and edge types, respectively. In this definition, each node v ∈ V is identified by a node type mapping ϕ(v) ∈ TV , and each edge e ∈ E is associated with an edge type mapping ψ(e) ∈ TE.

In addition to the structural aspects, nodes in a heterogeneous graph also have associated features. Different types of nodes often have different feature dimensions. Specifically, if a node v belongs to type t1, its feature vector has a dimension dt1 , while if a node u belongs to type t2, its feature vector has a dimension dt2 , where dt1 and dt2 are generally not equal.

4 Method

4.1 Overview Framework

In this paper, we aim to develop a graph-level pre-training task on a heterogeneous graph and learn a heterogeneous prompt vector capable of adjusting the original node features of the graph. By employing this heterogeneous prompt vector, we seek to bridge the gap between graph pre-training strategies and multiple downstream tasks, as well as to mitigate the challenges associated with transferring prior knowledge across different domains.To address this issue, we adopt a three-step approach. First, we unify all downstream tasks into a graph-level task format to minimize the gap between the pre-trained model and downstream task objectives. Next, recognizing that existing graph prompt learning methods struggle to effectively integrate contrastive pre-training strategies in the heterogeneous graph domain, we design a specialized graph-level contrastive pre-training strategy. This approach fully leverages heterogeneous information to enhance model performance in multi-task scenarios. Finally, we introduce heterogeneous feature prompts, which further improve model performance by optimizing the feature representation of input graphs. The overall framework of HGMF is shown in Fig 2.

4.2 Reformulating DownstreamTasks

Challenges in Heterogeneous Graphs and Task Reformulation

In homogeneous graphs, reformulating downstream tasks into a unified graph-level representation is a widely used approach to bridge the gap between pre-training and downstream tasks. By constructing τ-ego networks as induced subgraphs and mapping node or edge labels to these subgraphs, node and edge classification problems can be reformulated as graph classification tasks. However, in heterogeneous graphs, this method faces new challenges due to the diverse types of nodes and edges, each carrying different semantics. Additionally, the semantic dependencies between these types are complex and cannot be captured by simple τ-hop neighborhoods. The heterogeneity of these graphs makes it difficult to construct a consistent graph-level representation that can accommodate the full range of tasks.

Adapting Induced Subgraphs and Label Mapping

To address these challenges, we adapt the task reformulation method to heterogeneous graphs by modifying the construction of induced subgraphs. For node-level tasks, the induced subgraph includes the target node, its τ-hop neighbors, and all connecting edges along with their types, as show in Fig 1a.

(b) Induced subgraphs for edge

Figure 1: Induced subgraphs for nodes and edges

This subgraph captures both the structural context of the target node and the type-specific semantics of its neighboring nodes and edges. For edge-level tasks, the induced subgraph includes the two endpoints of the target edge and their respective τ-hop neighbors, retaining all nodes, edges, and their types within this subgraph, as depicted in Fig 1b. This ensures that the semantic context surrounding the target edge is preserved, which is critical for accurate edge-level task performance. After constructing the induced subgraphs, we map the labels of the downstream tasks to these subgraphs. For node-level tasks, the label of the target node is assigned to its induced subgraph, transforming node classification into a graph classification task. Similarly, for edge-level tasks, the label of the target edge (e.g., relationship type) is assigned to its induced subgraph, turning edge classification into a graph classification problem. This mapping unifies node-level and edge-level tasks under a consistent graph-level framework, enabling knowledge transfer across tasks.

Summary

By reformulating both node-level and edge-level tasks into graph classification tasks, we provide a unified framework suitable for heterogeneous graphs. This method preserves the important type-specific semantics of heterogeneous graphs and lays a solid foundation for pre-training and downstream prompt-based learning in such graphs.

4.3 Heterogeneous Graph Pre-training

Why Design Heterogeneous Graph Pre-training Tasks

To unify pre-training and downstream tasks at the graph level, thereby maximizing the task subspace and fully exploring the potential relationships in heterogeneous graphs, we propose to adopt graph-level pre-training strategies during the pretraining phase. However, there are currently no pre-training methods specifically designed for graph-level tasks in heterogeneous graphs. This limitation prevents existing strategies from effectively adapting to graph-level downstream tasks, thereby restricting their performance and generalization. Moreover, contrastive learning has demonstrated superior performance in representation learning across various domains by constructing positive and negative sample pairs to enhance feature discrimination. Therefore, we propose a novel graph-level contrastive pre-training strategy for heterogeneous graphs to address these gaps and leverage the advantages of contrastive learning.

How to Design Heterogeneous Graph Pre-training Tasks

Heterogeneous graphs contain diverse node and edge types that play a critical role in preserving their structural integrity. For instance, in the ACM dataset, "subject" nodes, despite comprising only 56 entities, connect to others through 3025 edges. To enhance model robustness without disrupting the graph structure, we propose two augmentation strategies: Heterogeneous Node Masking Augmentation and Heterogeneous Edge Permutation Augmentation, both guided by an augmentation ratio r. This ratio determines how much of the node or edge set is altered, ensuring variability while maintaining structural integrity.

For Heterogeneous Node Masking, the adjusted augmentation ratio for each node type is computed as:

where count(i) is the number of nodes of type i, and TV denotes all node types. Using this ratio, the final number of masked nodes is:

where a(i) is the adjusted ratio for node type i, r is the augmentation ratio, and |V | is the total number of nodes.

Similarly, for Heterogeneous Edge Permutation, the adjusted ratio for each edge type is:

where count(i) is the number of edges of type i, and TE represents all edge types. The number of permuted edges is given by:

where b(i) is the adjusted ratio for edge type i, and |E| is the total number of edges.

After applying the augmentation strategies, the processing of augmented graphs involves three key steps. First, a GNNbased Encoder is used, incorporating a feature transformation layer to preprocess input features and extract graph-level representations ZGi , ZGj . This flexible design ensures compatibility with various tasks. Second, a Projection Head maps high-dimensional representations into a smaller latent space using a two-layer MLP, which facilitates contrastive learning. Finally, the Contrastive Loss Function maximizes the similarity between positive pairs of augmented graphs while minimizing it for negative pairs. Given N graphs, 2N augmented graphs are generated per batch, and the loss is computed as:

where z+i and z-i are the positive and negative representations, and τ is the temperature coefficient. By optimizing this loss, the model learns robust and discriminative graph-level representations that effectively support downstream tasks.

Through this integrated approach, our augmentation process efficiently combines the unique characteristics of heterogeneous graphs with contrastive pre-training, offering a unified and powerful graph representation learning framework.

4.4 Heterogeneous Graph Prompt Feature (HGPF)

Why Design Heterogeneous Graph Prompt Feature (HGPF)

In heterogeneous graphs, due to the diversity of node types, the feature dimensions of each node type are typically different. Even when the feature dimensions are consistent across nodes, the meanings of these dimensions often vary between different node types. For example, in a scholarly paper graph, author nodes might represent their research domains, while paper nodes contain publication information, making the meaning of the same feature dimension vary across different node types.

To effectively adjust the input graph for better adaptation to pre-trained models during downstream tasks, a mechanism must be designed to modify the feature representations of nodes according to their types. Therefore, we introduce the Heterogeneous Graph Prompt Feature (HGPF) to provide distinct feature prompts for each node type. This allows the adjustment of features of different node types, improving their alignment with pre-training tasks.

How To Design Heterogeneous Graph Prompt Feature (HGPF)

For each node type, we design a specific feature prompt associated with that type. Within each node type, the same feature prompt is shared. The dimensionality of the prompt matches the dimensionality of the corresponding node type's features. The role of this prompt is to adjust the node's feature representation by adding or multiplying it with the node's original feature.

Let xi denote the feature vector of node i, where xi ∈ Rdi represents the feature vector of node i, and Ti is the node type. The corresponding feature prompt is pTi ∈ Rdi . The adjustment process for the node feature is as follows:

Through this process, the node features are optimized by the adjustment based on their original features and the feature prompt, making them more suitable for the pre-training task.

Why It Works?

The effectiveness of HGPF is based on the prompt mechanism, which works by adjusting the node features of the input graph to better align with the pre-training task. In graphbased tasks, the flexibility of feature space operations is crucial to the effectiveness of prompting. By designing feature prompts, we make node features more task-relevant during pre-training, facilitating effective knowledge transfer.

Fang et al [Fang et al., 2024]. demonstrated that for any graph G (with adjacency matrix A and node feature matrix X), an appropriate prompt token p∗ can be learned such that the following equation holds:

Here, g(A, X) represents a transformation applied to the original graph (e.g., changing node features, adding/removing edges), and Oφ denotes the error between the graph augmented by the prompt and the graph after transformation. This error depends on the model's nonlinear layers and the quality of the learned prompt. With further optimization, the error can be significantly reduced.

In our approach, HGPF adjusts node features by customizing task-specific prompts for each node type. Unlike traditional methods, our feature prompts are combined with the original features of the nodes through multiplication. This multiplication method enables a more rapid adaptation to few-shot scenarios, particularly when dealing with attributed graphs, thereby enhancing the node representations for task alignment. This adaptive adjustment mechanism facilitates a more effective transfer of knowledge from the pre-trained model to the downstream task, significantly improving the performance of downstream tasks.

5 Experiments

This section presents experimental results and analyses to demonstrate the superiority of our approach.

5.1 Experimental Setup

Datasets. We use two real-world HIN datasets, summarized in Table 1. (1) ACM [Lv et al., 2021] dataset is an academic network containing relationships among papers (P), authors (A), subjects (S), and terms (T), making it suitable for studying academic influence and knowledge discovery. (2) IMDB [Lv et al., 2021] dataset serves as a benchmark for graph neural networks, involving movies (M), actors (A), directors (D), and keywords (K), with a wide range of relationships reflecting its suitability for tasks in the entertainment domain.

Baselines. We evaluate our method against baselines spanning five categories: (1)Supervised GNNs: GCN [Kipf and Welling, 2016], GAT [Veliˇckovi´c et al., 2017]; (2) GNNs with "pre-train, fine-tune": GraphCL [You et al., 2020], SimGRACE [Xia et al., 2022]; (3) Supervised HGNNs: HGT [Hu et al., 2020], HAN [Wang et al., 2019]; (4) HGNNs with "pre-train, fine-tune": DMGI [Park et al., 2020], HeCo [Wang et al., 2021], HDMI [Jing et al., 2021]; (5) HGNNs with "pre-train, prompt": HGPROMPT [Yu et al., 2024].

Settings and parameters. For the few-shot setting, we evaluate tasks under a 10-shot configuration across three categories: node classification, edge classification, and graph classification. The task construction mainly follows the approach of ProG [Sun et al., 2023a] works. However, in heterogeneous graphs, target nodes are often connected to nontarget nodes, whose learning value is relatively lower. This makes directly applying the ProG [Sun et al., 2023a] approach, which combines both endpoints of edges as labels, less suitable in the heterogeneous domain.

To address this, we make specific adjustments: only the labels of nodes with endpoint types corresponding to target nodes are used as the induced graph's labels. The evaluation metrics include micro-F1 and macro-F1, ensuring a comprehensive assessment of the model's performance.For reproducibility, we provide the specific construction methods for each task and the detailed hyperparameter values in the supplement.

5.2 Performance Evaluation onMulti-Tasks

Few-shot node classification

The classification results of the 10-shot nodes are shown in Table 2. Based on the experimental results, we summarize our observations as follows:(1) HGMP outperforms all baseline methods across all datasets, demonstrating its effectiveness in node classification tasks under the few-shot setting within heterogeneous graph domains.(2) On the ACM dataset, HeCo, which also adopts a pre-training and finetuning strategy, achieves performance second only to HGMP and significantly surpasses homogeneous pre-training methods. This result highlights the necessity of fully leveraging heterogeneity information in pre-training methods for heterogeneous graphs. Furthermore, HeCo's performance, closely following that of HGMP, further validates the effectiveness of heterogeneous graph prompting mechanisms and the rationality of incorporating graph contrastive pre-training strategies that integrate heterogeneous information. Additionally, HAN, which utilizes meta-path strategies, also demonstrates competitive performance, showcasing its advantages in capturing structural and semantic information.(3) In fewshot scenarios, GCN significantly outperforms other homogeneous methods, indicating that models with smaller parameter sizes are more effective at avoiding overfitting when training data is limited, thereby achieving better generalization performance.

Few-shot edge classification

The results of the 10-shot edge classification task are shown in the Table 3. Based on the experimental results, we summarize the following observations: (1) HGMP outperforms all baseline methods in edge-type tasks, demonstrating its broad adaptability to few-shot settings in tasks within heterogeneous graph domains. (2) Since HeCo and DMGI only model the features of the target nodes, they fail to capture the complete information required for the edge classification task, thus losing their previous second-best performance in node classification tasks. GCN, once again, benefits from its smaller parameter size, achieving success in the few-shot setting due to its ability to avoid overfitting.

Few-shot graph classification

The results of the 10-shot graph classification task are shown in Table 4. Based on the experimental results, we summarize the following observations: (1) On the ACM dataset, all baseline methods improved significantly in graph-type tasks, due to the high concentration of same-label target nodes during subgraph construction. HGMP outperformed all baselines, demonstrating its superiority in multi-task few-shot settings within heterogeneous graphs. (2) The IMDB dataset, being a multi-label multi-class classification task, performed well in node and edge classification. However, in the graph classification task, to adhere to the previous task setup, we could only select one label, converting the problem into a singlelabel multi-class classification task. This conversion resulted in significant information loss, leading to a noticeable drop in performance across all baselines on the IMDB dataset in the graph classification task. (3) The strong performance of all baselines in the graph classification task highlights that the knowledge obtained from pretraining on graph-level tasks is better suited for downstream graph classification tasks. This further underscores the necessity of unifying task designs at the graph level. (4) GAT demonstrated a performance recovery, indicating that as the graph size increases slightly, attention-based models can regain competitive performance.

However, HeCo and DMGI, as they only model the features of target nodes, remain unable to capture complete graphlevel information and consequently fell out of the top three in this task.

5.3 Flexibility on backbones

To validate the adaptability and performance of the proposed HGMP model across different backbones, we conducted experiments using GCN, GAT, and HGT on node classification, edge classification, and graph classification tasks. HGMP consistently outperformed traditional supervised methods in node and edge classification. For GCN, HGMP achieved improvements of approximately 3% and 4% in node and edge classification, respectively, while maintaining comparable performance in graph classification. With GAT, HGMP excelled in edge classification, outperforming the baseline by nearly 1%, though it showed a slight performance decrease in graph classification. For HGT, HGMP provided noticeable gains in node and edge classification, albeit with modest improvements in graph classification. These results confirm HGMP's adaptability and effectiveness in few-shot heterogeneous graph tasks, leveraging graph augmentation and prompt features to enhance performance.

5.4 Performance with different shots

To analyze the impact of shot numbers, we conducted experiments on the ACM dataset for node classification (NC), edge classification (EC), and graph classification (GC). The results, shown in Figs 3, indicate that HGMP generally outperforms baselines, especially in EC and GC tasks. For NC, HGPROMPT excels in few-shot scenarios (fewer than 5 shots) due to its optimized class prototype but loses its edge as sample size increases, with HGMP showing stronger robustness and generalization. In EC and GC tasks, HGMP consistently achieves superior performance, attributed to its graph-level training objectives and stable adaptability across settings.

5.5 Ablation study

To evaluate the contribution of each component in HGMP, we conducted an ablation study comparing the full model with its variants, as shown in Table 5. All variants were trained using the reconstruction task template based on induced subgraphs to ensure fair comparisons. The results demonstrate that each component significantly enhances performance.

Comparing Variant 1 with Variant 3 highlights that the absence of heterogeneous graph augmentation causes performance degradation, validating the importance of optimizing augmentation strategies to capture complex structural information. Similarly, the comparison between Variant 1 and Variant 2 shows that heterogeneous prompt features provide substantial performance gains, effectively constructing taskrelevant representations and improving classification accuracy.

Overall, HGMP outperforms all variants, demonstrating the synergistic benefits of combining heterogeneous graph augmentation and prompt features. This design aligns induced subgraphs in downstream tasks with pretraining objectives, enhancing generalization and overall performance.

6 Conclusions

This paper proposes HGMP, a novel multi-task prompt framework that bridges pre-trained models and downstream tasks in heterogeneous graphs by unifying downstream tasks into a graph-level format and incorporating a graph-level contrastive pre-training strategy leveraging heterogeneous information. Experimental results demonstrate its significant advantages over existing methods, particularly in multi-task scenarios. While effective for attributed graphs, HGMP performs suboptimally on unattributed graphs, highlighting a need for improved adaptability. Future work will focus on addressing this limitation, offering new perspectives for heterogeneous graph learning and advancing multi-task learning and practical graph applications.

Acknowledgements

This work was supported in part by the Zhejiang Provincial Natural Science Foundation of China under Grant LDT23F01015F01, in part by the Key Technology Research and Development Program of the Zhejiang Province under Grant No. 2025C1023 and in part by the National Natural Science Foundation of China under Grants 62372146.