Source-linked AI summary
PRODIGY: Enabling In-context Learning Over Graphs
Qian Huang, Hongyu Ren, Peng Chen, Gregor Kržmanc, Daniel Zeng, Percy Liang, Jure Leskovec
TL;DR
Graph in-context learning—adapting to new tasks from examples without parameter updates—remains unexplored. PRODIGY uses prompt graphs and graph-based pretraining to support this setting, and reports improved performance over in-context baselines and limited-data finetuning.
Problem
How to enable graph models to perform diverse new tasks on novel graphs from prompt examples, without finetuning, remains unexplored.
Source: §1 · pp. 1–2Method
PRODIGY represents examples and queries with a prompt graph, then pretrains a graph neural network and attention mechanism using in-context objectives.
Source: §1 · p. 2Results
Across unseen citation-network and knowledge-graph tasks, PRODIGY exceeds contrastive pretraining with hard-coded adaptation by 18% on average and limited-data finetuning by 32.6% on average.
Source: §1 · p. 2Takeaways & Limitations
A PRODIGY-pretrained model can perform new classification tasks on new graphs through in-context learning, without retraining.
Source: §6 · p. 10The supervised pretraining objective requires extra labels when its target function is not part of the pretraining graph, and some graphs may lack such labels.
Source: §3.2.1 · p. 6
Abstract
from arXiv · showhide
In-context learning is the ability of a pretrained model to adapt to novel and diverse downstream tasks by conditioning on prompt examples, without optimizing any parameters. While large language models have demonstrated this ability, how in-context learning could be performed over graphs is unexplored. In this paper, we develop \textbf{Pr}etraining \textbf{O}ver \textbf{D}iverse \textbf{I}n-Context \textbf{G}raph S\textbf{y}stems (PRODIGY), the first pretraining framework that enables in-context learning over graphs. The key idea of our framework is to formulate in-context learning over graphs with a novel \emph{prompt graph} representation, which connects prompt examples and queries. We then propose a graph neural network architecture over the prompt graph and a corresponding family of in-context pretraining objectives. With PRODIGY, the pretrained model can directly perform novel downstream classification tasks on unseen graphs via in-context learning. We provide empirical evidence of the effectiveness of our framework by showcasing its strong in-context learning performance on tasks involving citation networks and knowledge graphs. Our approach outperforms the in-context learning accuracy of contrastive pretraining baselines with hard-coded adaptation by 18\% on average across all setups. Moreover, it also outperforms standard finetuning with limited data by 33\% on average with in-context learning.
1 Introduction
The paper addresses how to enable graph models to learn new tasks from examples while transferring across both tasks and graphs. It introduces prompt graphs and PRODIGY, then reports in-context performance gains on unseen citation networks and knowledge graphs.
1 Introduction
- Graph in-context learning remains unexplored, with the challenge of generalizing across both new tasks and new graphs without finetuning.Existing graph pretraining requires finetuning, while graph meta-learning targets different tasks within the same graph.Source: §1 · pp. 1–2
- The paper introduces prompt graphs as a unified representation connecting contextualized inputs, prompt examples, queries, and labels.The framework is intended to represent node-, edge-, and graph-level classification tasks in a form that can be given to the same model.Source: §1 · p. 2
- PRODIGY combines a graph neural network for node and edge representations with attention over prompt graphs and a family of in-context pretraining objectives.One objective, neighbor matching, classifies which neighborhood a node or edge belongs to.Source: §1 · p. 2
- PRODIGY improves over contrastive pretraining with hard-coded adaptation by 18% on average and over limited-data finetuning by 32.6% on average.The evaluation covers paper-category classification and knowledge-graph completion on graphs not used for training.Source: §1 · p. 2
2 In-context Learning over Graphs
This section defines graph classification tasks and few-shot prompting, then represents a prompt as a graph combining contextualized inputs with their examples, queries, and labels.
2 In-context Learning over Graphs
- The framework focuses on few-shot in-context learning for node and edge classification tasks over graphs.The section also notes that the formulation can extend to subgraph and graph-level classification.Source: §2 · p. 3, §2.1 · p. 3
2.1 Classification Tasks over Graphs
- Node classification predicts labels for individual graph nodes, while edge classification predicts labels for potential node pairs; the same formulation extends to subgraphs and graphs.For edge classification, labels may correspond to relations in a knowledge graph.Source: §2.1 · p. 3
2.2 Few-shot Prompting
- A few-shot prompt supplies labeled examples and query inputs, while the source graph provides context for those inputs.For a k-shot, m-way task, there are k examples per class; the model predicts labels for the queries.Source: §2.2 · p. 3
2.3 Prompt Graph Representation
- Each prompt input is contextualized with a data graph, formed by sampling its k-hop neighborhood from the source graph.The input node set is retained within this context graph.Source: §2.3 · p. 3
- The task graph connects data nodes for examples and queries to label nodes; known examples encode true and false labels, while query labels remain unknown.Together with the data graphs, this forms the prompt graph; the text also describes possible extensions to regression and free-form generation.Source: §2.3 · p. 4
3 Pretraining to Enable In-context Learning
PRODIGY assumes a pretraining graph independent of the downstream graph and represents few-shot classification tasks as prompt graphs, enabling in-context prediction without additional finetuning.
3.1 Message Passing Architecture over prompt graph
- A data-graph GNN embeds each input graph, then a task-graph GNN passes information among support examples, queries, and label nodes.Label representations are learned using support examples and propagated back to support and query representations; query logits use cosine similarity with label representations.Source: §3.1 · pp. 4–5
3.2 In-context Pretraining Objectives
- PRODIGY is designed to pretrain a model to solve diverse tasks in prompt-graph form with the same weights, enabling direct in-context use on downstream tasks.Its objectives are constructed in-context rather than training a shared encoder that needs a task-specific downstream head.Source: §3 · p. 4, §3.2 · p. 5
3.2.1 Pretraining Task Generation
- Pretraining objectives include self-supervised neighbor matching and supervised multi-task prompts, with the latter using available node- or edge-level labels.Neighbor matching classifies whether a node belongs to sampled local neighborhoods; supervised tasks can directly resemble downstream tasks but may require labels absent from the pretraining graph.Source: §3.2.1 · pp. 5–6
- Neighbor-matching tasks use sampled nodes as classes and their exact l-hop neighbors as support and query examples; the task can be adapted to classify edges for link prediction.Multi-task prompts can also be generated from node or edge labels, using edge types for link prediction.Source: §3.2.1 · pp. 5–6
3.2.2 Prompt graph generation with augmentation
- Prompt-graph augmentation randomly drops nodes or masks node features in each sampled k-hop subgraph before constructing the task graph.The stated motivation is to learn representations invariant to graph corruptions.Source: §3.2.2 · p. 7
3.2.3 Pretraining Loss
- The model is pretrained with cross-entropy objectives over prompt graphs generated for neighbor-matching and multi-task prompts.The loss uses model logits and corresponding query labels for the two task types.Source: §3.2.3 · p. 7
4 Experiments
PRODIGY is evaluated through in-context classification on citation and knowledge graphs, using few-shot prompts and comparisons with pretrained and finetuned baselines. It performs strongly across tasks, and the experiments examine component effects, prompt size, and scaling with pretraining data.
4.1 Experimental Setup
- Pretraining uses MAG240M and Wiki; evaluation covers arXiv, ConceptNet, FB15K-237, and NELL, with three-shot prompts sampled from limited labeled training examples.The study compares PG-NM, PG-MT, full PRODIGY, NoPretrain, Contrastive with nearest-neighbor adaptation, and Finetune with a linear head.Source: §4.1 · pp. 7–8
4.2 In-Context Learning Results
- 28.6% average and up to 48% improvement over Contrastive on arXiv, 12.2% average improvement over Contrastive on knowledge graphs, and 77.7% average improvement over Finetune on arXiv.The reported comparisons concern average accuracy across ways; the Finetune comparison requires additional downstream training.Source: §4.2 · p. 8
- PG-NM performs strongly on arXiv despite task mismatch, PG-MT adapts better on knowledge graphs but worse on arXiv, and full PRODIGY performs best across both settings.The authors suggest greater task diversity helps full PRODIGY avoid overfitting to a single strategy.Source: §4.2 · p. 8
- PG-NM achieves 69.07% on 3-way arXiv classification versus TENT’s 65.13%, despite PG-NM never seeing paper-category classification during pretraining.The comparison evaluates only test labels because TENT trains using arXiv training labels.Source: §4.2 · pp. 8–9
4.3 Ablations
- Removing attribute prediction from PG-NM causes an average 7% performance drop across ways, and using all objective elements yields the highest ablation performance.The results identify attribute prediction as the PG-NM component with the greatest impact.Source: §4.3 · p. 9
4.4 Evaluation using different numbers of in-context examples
- As prompt examples increase, the PG models’ margin over the baseline grows on ConceptNet, consistent with stronger use of task context.Figure 3 reports accuracy on ConceptNet in a four-way setting across prompt-example counts.Source: §4.4 · p. 9, Figure 3 · p. 9
4.5 Scaling with Data Size
- On arXiv in a five-way setting, PRODIGY improves with more pretraining data while Contrastive saturates quickly and fluctuates.Figure 4 illustrates the comparison as training steps and pretraining data increase.Source: §4.5 · p. 9, Figure 4 · p. 9
5 Related Work
The related work contrasts graph in-context learning with language-model prompting, conventional graph pretraining, and graph meta-learning. It emphasizes adaptation to new graphs and tasks without gradient updates.
5.1 In-context Learning of Large Language Models
- Language-model in-context learning uses prompt examples or textual instructions to adapt to diverse tasks without fine-tuning, a capability this work extends to graph data.The cited comparison with BERT describes reduced adaptation effort relative to fine-tuning.Source: §5.1 · p. 9
5.2 Pretraining on Graphs
- Earlier graph pretraining learns encoders for pretraining tasks but typically requires task-specific data and a classification head for downstream adaptation.PRODIGY instead targets direct in-context use without gradient updates.Source: §5.2 · p. 9
5.3 Meta Learning on Graphs
- Prior graph meta-learning generally trains and tests across tasks on the same graph, whereas this work targets new graphs and tasks without additional fine-tuning.The cited prior methods are reported to perform optimally when trained on similar curated tasks.Source: §5.3 · p. 10
6 Conclusion
The conclusion presents PRODIGY as a framework for applying new classification tasks to new graphs through prompt graphs, without downstream fine-tuning. The supplied appendix passages also identify an attribute-prediction loss and a dataset-statistics table.
6 Conclusion
- PRODIGY is presented as enabling new classification tasks on new graphs through prompt graphs, outperforming baselines including methods that use fine-tuning.The conclusion specifically names both node- and edge-classification tasks.Source: §6 · p. 10
A Attribute Prediction Loss
- Masked node features are reconstructed from learned embeddings using an MLP trained with mean squared error loss.This passage describes the attribute-prediction loss used with MaskNode augmentation.Source: A Attribute Prediction Loss · p. 12
B Dataset Statistics
- Table 3 is identified as presenting dataset statistics.Source: Table 3 · p. 12
C Task Graph GNN Architecture
The described architecture uses attention-based graph neural networks over task graphs, alongside data-graph message passing. The cited implementation passage distinguishes the Contrastive encoder from Finetune’s added classification head.
C Task Graph GNN Architecture
- The task-graph GNN uses attention, with each node attending to other nodes at every layer.Source: C Task Graph GNN Architecture · p. 12
D.1 Model Architecture, MAG240M and arxiv
- For citation-network PRODIGY and NoPretrain models, the implementation uses two DataGraph message-passing layers and one TaskGraph layer; Contrastive uses two DataGraph layers, while Finetune adds a linear head.The citation-network node features are initialized using a pretrained RoBERTa base model.Source: §D.1 · p. 12
D.2 Model Architecture, knowledge graph datasets
The knowledge-graph PromptGraph architecture uses separate message-passing layers for the DataGraph and TaskGraph. The MAG240M and Wiki pretraining setups use different batch sizes and neighbor-matching task frequencies, with evaluation on 500 test tasks each.
D.2 Model Architecture, knowledge graph datasets
- For knowledge-graph datasets, PromptGraph uses two DataGraph message-passing layers and two TaskGraph layers that pass messages along positive and query edges; Contrastive and Finetune use DataGraph layers and an aggregator.Finetune adds a linear layer head.Source: §D.2 · pp. 12–13
D.3 Training, MAG240M
- MAG240M pretraining uses 30-way, 3-shot tasks with 4 queries, batch size 1, and checkpoints every 500 steps; evaluation uses 500 test tasks with 3 shots and 3 queries.The setup uses an input dimension of 768, embedding dimension of 256, and AdamW with learning rate and weight decay of 1×10−3.Source: §D.3 · p. 13
D.4 Training, Wiki
- Wiki pretraining uses 30-way, 3-shot tasks with 4 queries, batch size 10, and checkpoints every 500 steps; evaluation uses 500 test tasks with 3 shots and 3 queries.Evaluation samples 1-hop neighborhoods for ConceptNet and FB15K-237, and 2-hop neighborhoods for NELL and Wiki.Source: §D.4 · p. 13
E Ablation on Table 4 for the PG-NM setting
The PG-NM ablation varies attribute prediction, DropNode and ZeroNode augmentations, and the number of shots. Attribute prediction has the largest reported impact on performance.
E Ablation on Table 4 for the PG-NM setting
- PG-NM ablations vary attribute prediction loss, DropNode and ZeroNode augmentations, with the default neighbor-matching task using 3 shots.The best setting combines all three elements.Source: E Ablation on Table 4 for the PG-NM… · pp. 13–14
- Removing attribute prediction causes an average 7% drop across all ways, the largest reported impact among the PG-NM ablations.Using all the ablated elements together gives the highest performance.Source: E Ablation on Table 4 for the PG-NM… · p. 14
- Removing DropNode and ZeroNode drops performance 3% on average; removing both attribute prediction and augmentations drops roughly 7%; reducing shots from 3 to 1 drops 3.5%, across all ways.Source: E Ablation on Table 4 for the PG-NM… · p. 14
F Evaluation using different numbers of shots
The figures track in-context learning accuracy as the number of prompt examples changes in 20-way classification settings on FB15K-237 and NELL.
F Evaluation using different numbers of shots
- Figure 5 tracks in-context learning accuracy on FB15K-237 across prompt-example counts in a 20-way setting.Source: Figure 5 · p. 14
- Figure 6 tracks in-context learning accuracy on NELL across prompt-example counts in a 20-way setting.Source: Figure 6 · p. 14
G Scaling with Data Size
As pretraining data increases, PRODIGY consistently improves on the arXiv 5-way setting, while the Contrastive baseline saturates quickly and fluctuates. The paper also discusses compute costs and potential social risks of relying on pretraining knowledge.
G Scaling with Data Size
- As train steps proxy more pretraining tasks, PRODIGY consistently improves while Contrastive saturates quickly and fluctuates; redundancy is 0.20% for 10k tasks.Source: G Scaling with Data Size · pp. 14–15
H Compute
- One 10k-step pretraining run takes 3 to 4 hours on one NVIDIA A100-SXM4-80GB GPU.Source: H Compute · p. 15
I Broader Impacts
- Overreliance on pretraining knowledge could increase social bias and unfair benefits to dominant groups; the paper recommends diverse, balanced pretraining data and downstream testing across groups and subdistributions.Source: I Broader Impacts · p. 15