PaperlayerView exact revision on arXiv

Source-linked AI summary

One for All: Towards Training One Graph Model for All Classification Tasks

Hao Liu, Jiarui Feng, Lecheng Kong, Ningyue Liang, Dacheng Tao, Yixin Chen, Muhan Zhang

arXiv:2310.00149v3cs.LG

TL;DR

Graph foundation models remain difficult to unify because graph domains use incompatible attributes and graph tasks require different representations. OFA addresses this with text-attributed graphs, LLM-based shared embeddings, nodes-of-interest, and graph prompting. A single OFA model performs well across domains and tasks in supervised, few-shot, and zero-shot settings, with particularly strong zero-shot results.

  • Problem

    Graph data differ across domains and graph tasks span node, link, and graph prediction, leaving unified graph modeling underexplored.

  • Method

    OFA describes nodes and edges with text, embeds them with an LLM, standardizes tasks using nodes-of-interest, and injects task information through graph prompting.

  • Results

    OFA performs well across domains and tasks in supervised, few-shot, and zero-shot evaluations, including zero-shot learning.

  • Takeaways & Limitations

    OFA provides a single graph model with in-context learning ability across different graph domains and classification task formats.

  • Takeaways & Limitations

    OFA currently cannot learn regression tasks, and its cross-domain datasets are limited.

Abstract

from arXiv · show

Designing a single model to address multiple tasks has been a long-standing objective in artificial intelligence. Recently, large language models have demonstrated exceptional capability in solving different tasks within the language domain. However, a unified model for various graph tasks remains underexplored, primarily due to the challenges unique to the graph learning domain. First, graph data from different areas carry distinct attributes and follow different distributions. Such discrepancy makes it hard to represent graphs in a single representation space. Second, tasks on graphs diversify into node, link, and graph tasks, requiring distinct embedding strategies. Finally, an appropriate graph prompting paradigm for in-context learning is unclear. We propose \textbf{One for All (OFA)}, the first general framework that can use a single graph model to address the above challenges. Specifically, OFA proposes text-attributed graphs to unify different graph data by describing nodes and edges with natural language and uses language models to encode the diverse and possibly cross-domain text attributes to feature vectors in the same embedding space. Furthermore, OFA introduces the concept of nodes-of-interest to standardize different tasks with a single task representation. For in-context learning on graphs, OFA introduces a novel graph prompting paradigm that appends prompting substructures to the input graph, which enables it to address varied tasks without fine-tuning. We train the OFA model using graph data from multiple domains (including citation networks, molecular graphs, knowledge graphs, etc.) simultaneously and evaluate its ability in supervised, few-shot, and zero-shot learning scenarios. OFA performs well across different tasks, making it the first general-purpose across-domains classification model on graphs.

1 INTRODUCTION

OFA addresses the difficulty of building a unified graph model across heterogeneous domains and task types by combining text-attributed graphs, LLM embeddings, nodes-of-interest, and graph prompting. Evaluations across supervised, few-shot, and zero-shot settings show that one model performs well across domains and tasks.

  • Approach: OFA uses graph prompting to inject task information into graph data, supporting in-context learning without fine-tuning model parameters.The prompting paradigm converts embedded inputs into prompted graphs with a unified task representation for adaptive prediction.
  • Motivation: Graph foundation models are difficult to unify because datasets differ substantially in feature representation, dimension, scale, and semantic meaning.Examples include molecular graphs with atom-feature index vectors and e-commerce graphs with Bag-of-Words item-description features.
  • Approach: OFA converts heterogeneous graph data into text-attributed graphs and embeds cross-domain text attributes into a shared space with an LLM.The framework describes graph nodes and edges using human-readable text before jointly learning across domains.
  • Approach: OFA introduces nodes-of-interest and an NOI prompt node to represent different graph task types through a common task-oriented structure.These components target the differing procedures required for node, link, and graph tasks.
  • Evaluation: A single OFA model performs well across domains and tasks under supervised, few-shot, and zero-shot evaluation, including zero-shot learning.The authors characterize zero-shot performance as particularly notable because most existing graph models cannot perform it.

2 PRELIMINARIES

The preliminaries define text-attributed graphs, classification learning scenarios, and language-model in-context learning. Together, these concepts establish the representations and evaluation settings used by OFA.

  • Text-attributed graphs: A text-attributed graph associates every node and edge with a text sentence and represents the graph as nodes, relations, and edges.OFA extends the concept to all graphs because nodes and edges can be described with text, including graphs whose raw features are not textual.
  • Learning scenarios: Supervised learning trains on labeled data and evaluates on validation and test sets whose labels are seen during training.The dataset is split into training, validation, and test subsets to select and assess the classifier.
  • Learning scenarios: Few-shot learning classifies queries using N-way K-shot support data, while zero-shot learning is the special case K = 0 with no support samples.Few-shot support data provide K labeled examples for each of N classes.
  • In-context learning: Language-model in-context learning uses demonstrations supplied through a prompt to condition the model’s generated task answer.The model receives prompt text containing relevant examples rather than requiring parameter fine-tuning for each task.

3 ONE-FOR-ALL: TOWARDS FOUNDATION MODEL ON GRAPH

OFA builds one graph foundation model by converting heterogeneous graphs into text-attributed graphs, standardizing task targets with Nodes-of-Interest, and injecting task information through graph prompting. It supports unified processing across node, link, and graph classification tasks, including in-context and zero-shot learning.

  • Unifying graph data: OFA converts graphs from multiple domains into text-attributed graphs, using natural-language descriptions for nodes and edges and LLM encodings in a shared feature space.The collected datasets span citation networks, molecular graphs, knowledge graphs, and other domains.
  • Evaluation and zero-shot learning: OFA is evaluated across nine datasets and supervised, few-shot, and zero-shot settings, with one model performing across domains and task types.The class-node design also permits arbitrary unseen class nodes with suitable text descriptions, supporting zero-shot prediction.
  • Unifying graph data: LLM-encoded node and edge text features provide fixed-length input vectors while retaining domain information for downstream graph learning.For node vi and edge eij, OFA defines representations as xi = LLM(svi) and xij = LLM(seij).
  • Unifying graph tasks: OFA represents node, link, and graph tasks through Nodes-of-Interest subgraphs and an NOI prompt node connected to the task’s target nodes.The NOI is one node for node classification, a node pair for link tasks, and all graph nodes for graph-level tasks.
  • Unifying graph tasks: The NOI prompt node combines task-description text with information from the target subgraph, enabling unified readout without task-specific pooling mechanisms.The resulting formulation turns different task levels into binary classification on class nodes so they can be trained together.
  • Graph prompting: Graph Prompting Paradigm appends prompt substructures containing NOI and class nodes to the input graph, allowing task information to be supplied through the graph itself.Class nodes connect bidirectionally to the NOI prompt node, while the prompted graph combines the input and prompt graphs.

4 RELATED WORKS

Graph prompting research has adapted prompt learning and subgraph-based approaches to unify graph tasks, while OFA extends these ideas toward cross-domain in-context learning.

  • Graph prompt-learning methods introduce trainable prompt vectors to extract task-relevant information for downstream graph tasks.
  • Subgraph-based approaches transform different graph tasks into graph-level tasks, providing an alternative route to task unification.

5 EXPERIMENTS

The experiments test whether text features and a single OFA model support supervised, few-shot, and zero-shot learning across graph domains and task types. Results indicate strong cross-domain supervised performance, domain-separated embeddings, and competitive low-resource generalization.

  • Experiment goals: The experiments evaluate text-feature effects, cross-domain versatility, language-model choices, and graph prompting for in-context learning.
  • Evaluation setup: The experiments use supervised results from Tables 2 and 3 and few-shot or zero-shot results from Tables 4–6 to assess OFA across task levels and datasets.
  • Cross-domain supervised learning: Supervised tests compare independently trained and jointly trained OFA models using four language models across nine graph datasets and multiple domains.
  • Cross-domain supervised learning: Independent and joint OFA training achieve comparable or better results than baseline methods on all evaluated datasets.
  • Cross-domain supervised learning: Joint OFA models perform well across domains, with OFA-st often outperforming OFA-ind-st and larger language models generally producing better, more stable joint-training performance.
  • Cross-domain supervised learning: NOI prompt-node embeddings from OFA-joint-st separate by domain in two-dimensional space, indicating domain-specific subspaces within the shared model representation.
  • Few-shot and zero-shot learning: Few-shot and zero-shot evaluations cover node, link, and graph tasks in transductive and transfer settings, with OFA showing comparable or better few-shot performance and zero-shot capability.

6 CONCLUSIONS, LIMITATIONS AND FUTURE RESEARCH

OFA combines language models and graph neural networks across diverse graph domains and task types, with processed datasets spanning citation, knowledge, wiki, and molecular graphs. The authors report strong supervised, few-shot, and zero-shot performance, while identifying regression and dataset breadth as future work.

  • Conclusions: OFA combines language models with graph neural networks to address graph learning across multiple application domains.The related work and implementation materials cover citation networks, molecular graphs, knowledge graphs, and wiki-based graphs.
  • Limitations and future research: OFA currently does not learn regression tasks, and the cross-domain datasets remain limited.The authors leave both extensions to future work.
  • Dataset coverage: The collected data include citation, knowledge, wiki, and molecular graph datasets with domain-specific textual node and edge features.Examples include paper titles and abstracts, Wikipedia content, entity descriptions and relation types, and molecule datasets.

B.2 VISUALIZATION OF OFA DATASET

The authors visualize OFA dataset embeddings by projecting sampled node representations into two dimensions. The visualization separates domains and distinguishes citation networks from different research areas.

  • Visualization procedure: The visualization randomly selects 400 node embeddings from each generated dataset and projects them to two dimensions using t-SNE.For molecular datasets, all node embeddings are included.
  • Visualization findings: Embeddings from molecular datasets, knowledge graphs, wiki pages, and citation networks are well separated in the visualization.The authors interpret this separation as evidence that the language model distinguishes the generated domain representations.
  • Visualization findings: Arxiv and Cora embeddings are close to each other but far from PubMed embeddings.Arxiv and Cora mainly contain computer-science papers, whereas PubMed focuses on biology.

C IMPLEMENTATION OF OFA

OFA standardizes task inputs as prompted text-attributed graphs, embeds their text with language models, and processes them with an edge-aware GNN. Attention aggregates representations across layers before class-node prediction.

  • OFA pipeline: All task inputs are standardized as prompted text-attributed graphs, with task information carried by prompt nodes and edges.This unified representation supports training across different task types.
  • Graph encoding: A GNN processes the prompted graph and embedded text through message-passing layers to obtain final node embeddings.The model extends relational graph convolution by incorporating edge features.
  • Representation aggregation: An attention layer aggregates outputs from all GNN layers to accommodate different task natures and mitigate over-smoothing.The attention can select important layers based on domain information in the text features.
  • Prediction: The model gathers class-node embeddings and uses an MLP prediction head for binary classification on each class node.For multi-class problems, OFA collects the class probabilities.
  • Experimental settings: Joint training uses seven GNN layers, while independent training uses six, with hidden size 768 and dropout 0.15.Initial embeddings from different language models are projected to 768 dimensions before the GNN.

D.2 FEW-SHOT AND ZERO-SHOT LEARNING EXPERIMENTS

The few-shot and zero-shot experiments evaluate OFA across transductive and transfer settings spanning node-, link-, and graph-level datasets. Few-shot prompting removes category information from class nodes so the model focuses on comparing query and support nodes.

  • Joint low-resource training: The low-resource joint-training experiments train one model on Arxiv node-level, FB15K237 link-level, and ChEMBL graph-level data.The experiments use limited label splits for Arxiv and FB15K237 and the whole ChEMBL dataset.
  • Few-shot prompting: Few-shot graph construction gives every class node a uniform text feature instead of category information.This shifts the task from learning class-specific information toward comparing query nodes with support nodes.
  • Evaluation settings: Testing covers transductive evaluation on ogbn-arxiv and FB15K237 and transfer evaluation on Cora, WN18RR, MOLPCBA, and MOLHIV.The benchmark spans node-level, link-level, and graph-level tasks.
  • Evaluation analysis: The study includes ablations of prompting design and few-shot accuracy tables for ogbn-arxiv, FB15K237, and WN18RR.The cited tables organize results by node-level and link-level tasks.

E.1 ABLATION STUDY

The ablation compares OFA’s graph prompting with a classification alternative that retains NOI subgraphs but removes the NOI prompt node. When datasets are trained jointly, OFA preserves performance while the alternative degrades.

  • Prompting ablation: The study compares OFA’s prompting paradigm with a “-Class node” alternative across HIV, ogbn-arxiv, Cora-node, and Cora-link datasets.The comparison covers joint and separate training settings.
  • Separate training: When datasets are trained separately, all methods achieve similar performance because prompting functions essentially as a pooling mechanism.
  • Joint training: When datasets are trained jointly, “-Class node” performance significantly drops, whereas OFA’s prompting approach maintains the original performance.
  • Joint training: The NOI prompt node carries task descriptions that help distinguish domains and tasks during joint training.Without it, the “-Class node” approach can confound tasks across domains.

E.2 FEW-SHOT AND ZERO-SHOT RESULTS

OFA is evaluated under joint and individual low-resource training across node-, link-, and graph-level tasks. Joint training helps some datasets, matches individual training on others, and can hurt link-level and some graph-level few-shot performance.

  • Experimental setup: The experiments compare one jointly trained low-resource model, OFA-joint-lr, with individually trained models for node-, link-, and graph-level tasks.The individual models use ogbn-arxiv, FB15K237, and Chemble for the respective task levels.
  • Node-level tasks: On Cora, OFA-joint-lr achieves better performance than OFA-ind-lr in all settings.The result may indicate that knowledge from link- and graph-level tasks helps Cora generalization.
  • Node-level tasks: On ogbn-arxiv, OFA-joint-lr and OFA-ind-lr achieve similar results.
  • Link-level tasks: For link-level tasks, OFA-ind-lr performs better than OFA-joint-lr, suggesting that knowledge from graph- or node-level tasks is not helpful.
  • Graph-level tasks: Joint training benefits HIV in most cases and PCBA in zero-shot settings, but PCBA few-shot performance drops significantly.The authors suggest differing numbers of training tasks as one possible reason.

F LIMITATIONS AND FUTURE WORKS

The paper identifies limitations in regression coverage and training-data scale. It leaves range-aware regression and stronger unsupervised training techniques for future work.

  • Limitations: OFA cannot yet perform regression tasks because regression targets can be unbounded in value.The paper therefore focuses on general classification.
  • Limitations: If zero-shot regression targets fall outside the training range, OFA may struggle to predict the correct target value.
  • Future work: Specifying a target range in the NOI prompt is proposed as a possible regression approach, but mathematical reasoning is considered unreliable with the current language models.The approach is left for future work.
  • Limitations: OFA’s graph-foundation-model training data remains scarce compared with the training data available for LLMs.
  • Future work: Auto-regressive, contrastive, and other unsupervised training techniques are identified as possible ways to further enhance OFA’s performance.The paper treats these techniques as an important future direction.
Loading 2310.00149v3…