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

Authors: Zehong Wang, Zheyuan Zhang, Tianyi Ma, Nitesh V Chawla, Chuxu Zhang, Yanfang Ye
Primary area: foundation or frontier models, including LLMs
Keywords: Graph Neural Network, Graph Foundation Model, Transferable Pattern, Tree Structure

Abstract:
Foundation models aim to create general, cross-task, and cross-domain machine learning models by pretraining on large-scale datasets to capture shared patterns or concepts (generalities), such as contours, colors, textures, and edges in images, or tokens, words, and sentences in text. However, discovering generalities across graphs remains challenging, which has hindered the development of graph foundation models. To tackle this challenge, in this paper, we propose a novel approach to learn generalities across graphs via task-trees. Specifically, we first define the basic learning instances in graphs as task-trees and assume that the generalities shared across graphs are, at least partially, preserved in the task-trees of the given graphs. To validate the assumption, we first perform a theoretical analysis of task-trees in terms of stability, transferability, and generalization. We find that if a graph neural network (GNN) model is pretrained on diverse task-trees through a reconstruction task, it can learn sufficient transferable knowledge for downstream tasks using an appropriate set of fine-tuning samples. To empirically validate the assumption, we further instantiate the theorems by developing a cross-task, cross-domain graph foundation model named Graph generality Identifier on task-Trees (GIT). The extensive experiments over 30 graphs from five domains demonstrate the effectiveness of GIT in fine-tuning, in-context learning, and zero-shot learning scenarios. Particularly, the general GIT model pretrained on large-scale datasets can be quickly adapted to specific domains, matching or even surpassing expert models designed for those domains.

Decision: reject
Reviewer scores: 6, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1730010790352

thread_type:
official_review

<thread_object>
creator: reviewer_bmnx

<content>
summary:
The paper centers on learning shared structures, or “generalities,” across various graphs using a unique construct called "task-trees." The proposed model, Graph Generality Identifier on task-Trees (GIT), is designed for broad applicability across graph-related tasks, such as node, edge, and graph classification.

strengths:
1. The presentation is generally clear.
2. Both theoretical and experimental sections are well-executed, with experiments spanning a range of graph-based tasks across multiple domains.

weaknesses:
1. My primary question is: what exactly differentiates the proposed Task-Tree Space from existing concepts like prompt graphs, contextual graphs, anchor nodes, or nodes of interest[1][2][3][4]? We already know that subgraphs can be used to unify node-, edge-, and graph-level tasks, and this paper seems to be doing the same thing, just under a different name.
2. While task-trees provide a generalized approach, there is a potential trade-off in terms of performance on domain-specific tasks, as noted in the specialization experiments. 
3. Could the authors clarify the computational impact of task-tree transformations, particularly for large-scale graphs?

[1]Sun, X., Cheng, H., Li, J., Liu, B., & Guan, J. (2023, August). All in one: Multi-task prompting for graph neural networks. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (pp. 2120-2131).

[2]Liu, H., Feng, J., Kong, L., Liang, N., Tao, D., Chen, Y., & Zhang, M. (2023). One for all: Towards training one graph model for all classification tasks. arXiv preprint arXiv:2310.00149.

[3]He Y, Hooi B. UniGraph: Learning a Cross-Domain Graph Foundation Model From Natural Language[J]. arXiv preprint arXiv:2402.13630, 2024.

[4]Huang, Q., Ren, H., Chen, P., Kržmanc, G., Zeng, D., Liang, P. S., & Leskovec, J. (2024). Prodigy: Enabling in-context learning over graphs. Advances in Neural Information Processing Systems, 36.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730260978771

thread_type:
official_review

<thread_object>
creator: reviewer_htag

<content>
summary:
The paper presents GIT, learning general patterns using task-trees. By pretraining on task-trees, GIT captures transferable knowledge, validated through experiments on 30 graphs across five domains. GIT demonstrates strong performance in fine-tuning, in-context, and zero-shot learning, often outperforming domain-specific models.

strengths:
1. The paper is well-written and easy to follow.
2. The paper provides a theoretical analysis of the effectiveness of task-trees.

weaknesses:
1. **Inaccurate Statements**: The authors make some inaccurate statements. For example, in Appendix A, they claim that "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." However, many foundation models [1,2] do not rely on subgraphs.
2. **Limited Novelty in Task Trees as Unified Instances**: The novelty of leveraging task trees as unified instances is limited. Building task trees based on task-relevant nodes as unified instances is quite similar to using subgraphs as instances [3,4]. Subgraphs are a versatile structure that can represent node-, edge-, and graph-level instances [3].
3. **Limitations in Building a Graph Foundation Model**: I believe the paper is far from building a graph foundation model. Although a unified instance is proposed to unify various tasks, the semantic and structural gaps between various domains [1,2] remain unaddressed. These are important challenges in building a foundation model.

[1] Zhao et al. "All in one and one for all: A simple yet effective method towards cross-domain graph pretraining." Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2024.\
[2] Yu et al. "Text-Free Multi-domain Graph Pre-training: Toward Graph Foundation Models." arXiv preprint arXiv:2405.13934 (2024).\
[3] Liu et al. Graphprompt: Unifying pre-training and downstream tasks for graph neural networks[C]//Proceedings of the ACM Web Conference 2023. 2023: 417-428.\
[4] Hao et al. Motif-based prompt learning for universal cross-domain recommendation[C]//Proceedings of the 17th ACM International Conference on Web Search and Data Mining. 2024: 257-265.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730654862460

thread_type:
official_review

<thread_object>
creator: reviewer_wqqf

<content>
summary:
This paper proposes GIT, which learns generalities across graphs via a novel concept of task trees. The task trees are defined as the basic learning instances in graphs and they can preserve the generalities shared across graphs. The extensive experiments demosntrate the effectiveness of GIT in fine-tuning, in-context learning, and zero-shot learning.

strengths:
1. This paper proposes a novel concept of task trees.
2. The authors show that task trees can preserve the generalities shared across graphs by several theorems and extensive experiments.

weaknesses:
1. As pointed out by Line 90, message-passing GNNs---which have been used in many existing graph prompt methods---can fully capture task trees. Thus, the existing graph prompt methods may also implicitly use task trees (the definition may be different). Please compare GIT with the existing graph prompt methods in detail and explain why GIT can capture the task generalities better than the existing graph prompt methods.
2. I suggest moving the related work into the main body.

questions:
See Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730784481638

thread_type:
official_review

<thread_object>
creator: reviewer_n4ou

<content>
summary:
This paper presents a solution towards graph foundation model which works in cross-task and cross-domain scenarios. With a task-tree-based appraoch to unify different graph tasks, along with a reconstruction-based pretraining objective, the authors provide theoretical analysis on the stability and transferability of their proposed method. By pretraining on a diverse set of graph datasets, the pretrained model is able to generalize to unseen graphs from various domains, and the specification process can further improve the performance.

strengths:
1. The paper is well written and contains extensive information from theoretical and empirical perspectives. 

2. Based on the experimental results, the proposed method achieves the goal of 'cross-task, cross-domain' transfer learning, which is quite impressive compared with existing GFM research.

weaknesses:
I would urge the authors to provide a more detailed discussion of the difference between their proposed 'task-tree' method and existing subgraph-based methods.

questions:
1. What are the differences between the proposed 'task-tree' method and existing subgraph-based methods? Could you explain this further, in terms of efficiency (i.e., why the task-tree method is more efficient), and effectiveness (i.e., how can the task-tree representation improve the performance of graph tasks). 

2. The authors adopted a reconstruction-based objective for pretraining. Could you provide some intuitions to justify this choice? Is it possible to use a contrastive learning method instead? 

3. In the fine-tuning experiments, did the model go through the specification process, or were they directly fine-tuned on specific downstream datasets after general pretraining?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731882091340

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We sincerely thank the reviewers for their valuable feedback and constructive suggestions.

We would like to take this opportunity to reiterate our contributions: In the context of graph foundation models, identifying general patterns across graphs (like shared tokens in languages) from diverse domains is essential. **To the best of our knowledge, this work is the first to theoretically demonstrate the existence of a graph learning instance, which we term the "task-tree", that preserves these cross-domain generalities.** 

Some reviewers raised concerns that our proposed task-tree is similar to subgraphs (or equivalent concepts like prompt graphs, contextual graphs, NOI graphs, etc., found in existing literature), and therefore questioned the novelty of our approach. However, **our proposed task-tree concept is fundamentally different from subgraphs**, as we will clarify below.

> 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). 

| Batch Size | 512 | 1024 | 2048 | 4096 | 8192 | 16384 | 32768 |
| --- | --- | --- | --- | --- | --- | --- | --- |
| *GIT - Task-Tree* |  |  |  |  |  |  |  |
| Time / Epoch (s) | 208 | 180 | 176 | 172 | 172 | 163 | 162 |
| Memory Allocation (%) | 6 | 8 | 18 | 41 | 75 | 93 | 98 |
| *GIT - SubG* |  |  |  |  |  |  |  |
| Time / Epoch (s) | 280 | 243 | 234 | 223 | OOM | OOM | OOM |
| Memory Allocation (%) | 21 | 39 | 74 | 97 | OOM | OOM | OOM |

**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. 

|  | OFA | GIT - SubG | **GIT** |
| --- | --- | --- | --- |
| Academia (node classification) | 72.18 | 73.48 | **75.82** |
| KG (edge classification) | 72.38 | 73.59 | **75.73** |
| Molecule (graph classification) | 74.03 | 72.67 | **75.73** |

---

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.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The submission proposes a task-tree framework for graph foundation models, which aims to unify node, edge, and graph-level tasks. The reviewers have raised concerns regarding the novelty of the proposed task-tree concept, its differentiation from existing methods, and the effectiveness of the theoretical claims. While the reviewers acknowledged the paper's strengths, no reviewer strongly championed acceptance, and one reviewer explicitly recommended rejection; thus, I lean towards rejecting this submission.