Towards Graph Foundation Models: Learning Generalities Across Graphs via Task-trees

Dear Reviewer BMnX,

We sincerely appreciate your detailed feedback and constructive criticism, which will help us improve our work. Below, we address your concerns in detail:

W1: Differentiation of Task-Tree Space from Existing Concepts

Thank you for highlighting this concern. We would like to clarify that all the referenced works [1,2,3,4] are fundamentally based on subgraphs, whereas our proposed approach leverages task-trees. We will provide a thorough discussion on the distinctions between subgraphs and task-trees. In summary, our task-tree approach is both more efficient and effective in unifying graph-related tasks and preserving generalities across different graphs.

The indicated methods are based on subgraphs: Here is a more detailed breakdown of why these methods are subgraph-based:

• [1]: This work uses graph prompt learning for handling different graph-related tasks by extracting subgraphs around task-relevant nodes (e.g., ego-graphs around nodes for node classification). GNNs are then applied to the induced subgraphs to obtain embeddings for predictions.
• [2]: This method enables in-context learning for node classification using what it refers to as "data graphs". However, these are essentially k-hop ego-graphs extracted around each node, which are then integrated into a graph prompt learning framework to learn subgraph embeddings.
• [3]: Building on [1] and [2], this work extends the framework to a cross-domain, cross-task graph foundation model. Here, it uses the concept of NOI (Nodes of Interest) graphs to unify graph tasks, but the process is fundamentally the same as extracting subgraphs.
• [4]: This work builds upon the subgraph sampling methods of [1], employing k-hop ego-graphs and using techniques like personalized PageRank to select subgraphs more effectively.

Following, we will present the detailed comparison between subgraphs and task-trees.

Subgraphs: Some existing methods utilize subgraphs (typically k-hop subgraphs) as the basic learning instances [1,2,3]. For example, in node classification, these approaches extract ego-graphs around each node and assign the labels of the central nodes as labels for the induced subgraphs, transforming node-level classification into subgraph-level classification. This approach can also be adapted for link-level or graph-level tasks [1]. In general, these methods involve: (1) Extracting subgraphs around task-relevant nodes [1] and (2) Applying GNNs on each subgraph to obtain subgraph embeddings for predictions. However, these methods have two significant limitations: (1) Efficiency Issue: Extracting subgraphs incurs additional computational overhead, leading to increased time and memory usage due to the need to process and store these subgraphs. (2) Learnability Issue: The information preserved in subgraphs may not always be learnable by standard GNNs. Given the expressiveness limitations of message passing (bounded by the 1-WL test), these methods often struggle to capture essential substructures within subgraphs.

Task-Trees: Compared to subgraphs, task-trees are both more efficient and learnable. The encoding of task-trees for a node/link/graph involves: (1) Adding virtual nodes to the original graphs and connecting them to task-relevant nodes. (2) Using GNNs over the augmented graph to encode the embeddings of these virtual nodes for predictions. For instance, in node classification, we first augment the original graph by adding virtual nodes connected to each node and perform GNNs on the new graph to learn virtual node embeddings for predictions. Thus, our method requires only the addition of nodes and edges to the original graph, making it significantly more efficient than extracting and storing the subgraph. Additionally, encoding task-trees is equivalent to directly encoding the virtual nodes via message passing, ensuring that the information in task-trees is fully learnable by standard GNNs. In conclusion, our task-tree approach offers both superior efficiency and learnability.

Dear Reviewer hTAG,

We sincerely appreciate your thorough review and constructive feedback, which will help us enhance our work. Below, we address your concerns in detail:

W1: Inaccurate Statements

Thank you for pointing this out. We apologize for any confusion and would like to clarify the statement in question. In Appendix A, we stated: “… Other efforts to resolve feature heterogeneity include methods like singular vector decomposition (SVD) [1,2] and non-parametric encoders [4]. However, most of these approaches rely on subgraphs as the primary learning instances, which can result in inefficient training and reduced expressiveness, as discussed in the main paper.”

Our intention was to convey that these methods [1,2,4] can handle feature heterogeneity; however, we did not mean to imply that [1,2,4] are based on subgraphs for task unification. In fact, these methods do not explicitly address the unification of node, link, and graph tasks but instead directly apply GNNs to learn node embeddings, which are subsequently used for model training. We acknowledge that the wording may have been misleading, and we have revised this section in the updated version to avoid any further misunderstandings.

W2: Limited Novelty in Task Trees as Unified Instances

Regarding your concern about the novelty of task-trees as unified instances, we’d like to provide a detailed comparison between subgraphs and task-trees. In summary, our proposed task-trees are more efficient and learnable compared to subgraphs, particularly in unifying various graph-related tasks and preserving generalities across different domains.

Subgraphs: Some existing methods utilize subgraphs (typically k-hop subgraphs) as the basic learning instances [3,5,6]. For example, in node classification, these approaches extract ego-graphs around each node and assign the labels of the central nodes as labels for the induced subgraphs, transforming node-level classification into subgraph-level classification. This approach can also be adapted for link-level or graph-level tasks [5]. In general, these methods involve: (1) Extracting subgraphs around task-relevant nodes [5] and (2) Applying GNNs on each subgraph to obtain subgraph embeddings for predictions. However, these methods have two significant limitations: (1) Efficiency Issue: Extracting subgraphs incurs additional computational overhead, leading to increased time and memory usage due to the need to process and store these subgraphs. (2) Learnability Issue: The information preserved in subgraphs may not always be learnable by standard GNNs. Given the expressiveness limitations of message passing (bounded by the 1-WL test), these methods often struggle to capture essential substructures within subgraphs.

Task-Trees: Compared to subgraphs, task-trees are both more efficient and learnable. The encoding of task-trees for a node/link/graph involves: (1) Adding virtual nodes to the original graphs and connecting them to task-relevant nodes. (2) Using GNNs over the augmented graph to encode the embeddings of these virtual nodes for predictions. For instance, in node classification, we first augment the original graph by adding virtual nodes connected to each node and perform GNNs on the new graph to learn virtual node embeddings for predictions. Thus, our method requires only the addition of nodes and edges to the original graph, making it significantly more efficient than extracting and storing the subgraph. Additionally, encoding task-trees is equivalent to directly encoding the virtual nodes via message passing, ensuring that the information in task-trees is fully learnable by standard GNNs. In conclusion, our task-tree approach offers both superior efficiency and learnability.

Efficiency Comparison (Section 5.4 of the Original Submission): We conducted an empirical comparison between subgraphs and task-trees. For this, we implemented a subgraph version of our model, GIT-SubG, by replacing task-trees with subgraphs. Below is a comparison of time and memory usage (on a 48GB GPU) between GIT and GIT-SubG during pretraining (about 1,700,000 task-trees). To give a better understanding, using a batch size of 2048, GraphMAE requires 193 seconds per epoch with a memory allocation of 35%.

Dear Reviewer WQqF,

We sincerely appreciate your recognition of our contributions and your insightful feedback, which will help us further improve our work. Below, we address your concerns in detail:

W1: Comparison Between GIT and Existing Graph Prompt Methods

Thank you for pointing this out. We acknowledge that some concepts in existing graph prompt learning methods, such as induced graphs [1], data graphs [2], or NOI graphs [3], may appear to have similarities to our proposed task-trees. However, these methods fundamentally rely on subgraphs, not task-trees. In particular, these approaches involve: Extracting subgraphs (typically k-hop ego-graphs) around task-relevant nodes. Applying GNNs and graph prompt learning to each induced subgraph to learn subgraph representations, which are then used for predictions. For example, in node classification, existing methods extract ego-graphs around each node and assign labels to the induced subgraphs based on the central node, effectively transforming node-level classification into subgraph-level classification. However, subgraph-based methods have two inherent limitations: (1) Efficiency Issue: Extracting subgraphs introduces significant computational overhead, increasing both time and memory usage due to the need to process and store numerous subgraphs. (2) Learnability Issue: The information preserved in subgraphs may not always be fully captured by standard GNNs. Due to the expressiveness limitations of message passing (restricted by the 1-WL test), these methods often struggle to capture essential substructures within the subgraphs.

Our proposed task-tree method fundamentally differs from the subgraph-based approaches. Task-trees offer both higher efficiency and improved learnability. The encoding process for task-trees involves: (1) Virtual node augmentation: Adding virtual nodes to the original graph and connecting them to task-relevant nodes. (2) GNN encoding: Using GNNs on the augmented graph to generate embeddings for these virtual nodes, which are then used for predictions. For instance, in node classification, we augment the graph by adding virtual nodes linked to each task-relevant node and use GNNs to learn embeddings for these virtual nodes. This approach is significantly more efficient since it only requires augmenting the original graph without extracting subgraphs. Moreover, since encoding task-trees is equivalent to encoding virtual nodes through message passing, the task-tree structure is fully learnable by standard GNNs.

Conceptually, task-trees are similar to computation trees used in GFT [4], which defines transferable vocabularies for graph foundation models. However, our work goes further by providing a theoretical foundation that explains why tree structures (such as our task-trees and the computation trees in [4]) can serve as fundamental learning instances for graphs.

Effectiveness Comparison (Section 5.4 of the Original Submission): We conducted an empirical comparison between subgraphs and task-trees. For this, we implemented a subgraph version of our model, GIT-SubG, by replacing task-trees with subgraphs. We evaluated model performance across three domains: academic networks (node classification), knowledge graphs (edge classification), and molecule graphs (graph classification). Additionally, we compared our model against another popular subgraph-based method, OFA [3], which employs graph prompt learning for downstream tasks. The results (averaged over all graphs in each domain) are as follows.

W2: Move Related Work into Main Body

Thank you for the suggestion. We will incorporate a section on related work into the main paper.

Thank you once again for your valuable feedback. Please let us know if you have any further questions or need additional clarification.

Reference:

[1] All in One: Multi-task Prompting for Graph Neural Networks, KDD 23.

[2] PRODIGY: Enabling In-context Learning Over Graphs, NeurIPS 23.

[3] One for All: Towards Training One Graph Model for All Classification Tasks, ICLR 24.

[4] GFT: Graph Foundation Model with Transferable Tree Vocabulary, NeurIPS 24.

Comparison to Subgraphs — A Case Study: Our experiments (detailed below) indicate that subgraph encoding is the primary bottleneck. For instance, consider a graph with 100 nodes and 300 edges (average degree = 3). Sampling 2-hop ego-graphs for all nodes results in subgraphs with ~10 nodes and ~9 edges each. To encode these subgraphs simultaneously, they must be combined into a larger graph with 1,000 nodes and 900 edges. Then, they should conduct message passing on the larger graph to learn subgraph embeddings. By contrast, our task-tree approach simply adds 100 virtual nodes and 100 edges to the original graph, resulting in a graph with 200 nodes and 400 edges—making the message passing more efficient

W4: A motivating example

To motivate the concept of task-trees, we begin by considering the structures shared across different graph tasks. Existing subgraph-based methods typically extract subgraphs around task-relevant nodes to form induced graphs, unifying all graph tasks as graph classification. However, we approach this problem from a different perspective, focusing on the learning process of GNNs. Regardless of the graph task, predictions ultimately rely on the embeddings of task-relevant nodes. For instance, in graph classification, embeddings of all nodes in the graph are aggregated for prediction. Existing methods achieve this aggregation using a separate pooling module, but we ask: why not simply append virtual nodes to the original graphs and use message passing to directly mimic this aggregation process? By doing so, predictions for any graph-related task can be made directly through the embeddings of these virtual nodes.

By adding virtual nodes to the original graph, we construct what we call task-trees, which stem from these virtual nodes. Conceptually, while existing methods unify graph tasks as graph classification, our approach unifies these tasks as node classification, specifically task-tree classification. This distinction simplifies the problem and allows us to leverage the natural compatibility of message passing with tree-like structures.

Regarding the performance of task-trees, we hope our previous responses have provided sufficient case studies to highlight their advantages in terms of learnability and efficiency compared to subgraphs.

We appreciate your thoughtful consideration and remain available to provide any further clarifications or additional insights as needed.

Dear Reviewer N4ou,

We sincerely appreciate your recognition of our contributions and your insightful feedback, which will help us further refine our work. Below, we address your concerns in detail:

W1: Difference between Task-Trees and Subgraphs

Subgraphs: Some existing methods utilize subgraphs (typically k-hop subgraphs) as the basic learning instances [1,2,3]. For example, in node classification, these approaches extract ego-graphs around each node and assign the labels of the central nodes as labels for the induced subgraphs, transforming node-level classification into subgraph-level classification. This approach can also be adapted for link-level or graph-level tasks [1]. In general, these methods involve: (1) Extracting subgraphs around task-relevant nodes [1] and (2) Applying GNNs on each subgraph to obtain subgraph embeddings for predictions. However, these methods have two significant limitations: (1) Efficiency Issue: Extracting subgraphs incurs additional computational overhead, leading to increased time and memory usage due to the need to process and store these subgraphs. (2) Learnability Issue: The information preserved in subgraphs may not always be learnable by standard GNNs. Given the expressiveness limitations of message passing (bounded by the 1-WL test), these methods often struggle to capture essential substructures within subgraphs.

Task-Trees: Compared to subgraphs, task-trees are both more efficient and learnable. The encoding of task-trees for a node/link/graph involves: (1) Adding virtual nodes to the original graphs and connecting them to task-relevant nodes. (2) Using GNNs over the augmented graph to encode the embeddings of these virtual nodes for predictions. For instance, in node classification, we first augment the original graph by adding virtual nodes connected to each node and perform GNNs on the new graph to learn virtual node embeddings for predictions. Thus, our method requires only the addition of nodes and edges to the original graph, making it significantly more efficient than extracting and storing the subgraph. Additionally, encoding task-trees is equivalent to directly encoding the virtual nodes via message passing, ensuring that the information in task-trees is fully learnable by standard GNNs. In conclusion, our task-tree approach offers both superior efficiency and learnability.

Efficiency Comparison (Section 5.4 of the Original Submission): We conducted an empirical comparison between subgraphs and task-trees. For this, we implemented a subgraph version of our model, GIT-SubG, by replacing task-trees with subgraphs. Below is a comparison of time and memory usage (on a 48GB GPU) between GIT and GIT-SubG during pretraining (about 1,700,000 task-trees). To give a better understanding, using a batch size of 2048, GraphMAE requires 193 seconds per epoch with a memory allocation of 35%.

Effectiveness Comparison (Section 5.4 of the Original Submission): We also evaluated model performance across three domains: academic networks (node classification), knowledge graphs (edge classification), and molecule graphs (graph classification). Additionally, we compared our model against another popular subgraph-based method, OFA [3], which employs graph prompt learning for downstream tasks. The results (averaged over all graphs in each domain) are as follows.