PRODIGY: Enabling In-context Learning Over Graphs
Authors: Qian Huang
Hongyu Ren
Peng Chen
Gregor Kržmanc
Daniel Zeng
Percy Liang
Jure Leskovec
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Graph Neural Network
in-context learning
pretraining
TL;DR: We enable in-context learning over graphs with a novel in-context task representation and corresponding pretraining architecture and objectives.

Abstract: 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.

PRODIGY: Enabling In-context Learning Over Graphs
Qian Huang'*
Hongyu Ren hyrenccs stanford.
ghwang@cs. stanford.
Peng Chen' pengo@stanford.co edu
Gregor Kržmanc² gregor.krzmance krzzancciijaaj
Daniel Zeng'
dzengocs stanford.
Percy Liang?
Jure Leskovec jureêcs setreccsstaanfoo coanford.com edu
pliangucs stanford.edu
1 Stanford University
2 University of Ljubljana
Abstract
In-context learning is the ability pretrained model to adapt to novel and diverse downstream tasks by conditioning 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 Pretraining Over Diverse In-Context Graph Systems (PRODIGY), the first pretraining framework that enables in-context learning over graphs. The key idea of our framework to formulate in-context learning over graphs with a novel 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 pre- trained 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% average across all setups. Moreover, it also outperforms standard finetuning with limited data by 33% on average with in-context learning.
1 Introduction
In-context learning is a novel and one of the most intriguing capabilities of language models [1]. referss the capability a pretrained model perform novel and diverse tasks directly at the prediction time when prompted with just a few examples, without the need to update the model weights. For example, a person may describe the new task (e.g., question answering. machine translation, or code generation) using natural language and demonstrate it to the language model with several prompt examples. The language model then directly without any model training or finetunning performs the task.
However, how to enable in-context learning for diverse graph machine learning tasks, such as identifying misinformation spreader in social networks [14] and product suggestions across online e-commerce websites [21], still remain unexplored and challenging. An in-context learner for graphs should be able solve novel tasks on novel graphs. For example, give music product
indicates equal contribution.
37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

Few-shot Prompting
Prompt Graph
(C)
(A)
Graph g:
Prompt Examples 5 Input
Task Graph
[0,
Labely
Fg
go
Queries
Data Graphs
Prompt examples S
Queries
00000 input nodes (head)
Input nodes (tail)
BHCE data nodes
oo label nodes
Figure 1: In-context few-shot prompting over graphs with prompt graph for edge classification in PRODIGY. (A) Given the source graph 9, we provide prompt examples S that consist of the input head/tail nodes and their labels, as well as the queries. (B) For each datapoint from both prompt examples and the queries, we first construct its data graph Go by retrieving context from the source graph G. (C) Then we create a task graph to capture the connection between each datapoint and each label, which includes a data node Us for each datapoint and label node Uy for each label in V. Each pair of data and label nodes are connected with edge attributes corresponding to their binary labels.
recommendations on Spotify when being trained on Amazon purchasing graph. The first challenge here is how to formulate and represent node-, edge- and graph-leve tasks over graphs with a unified task representation that allows the model to solve diverse tasks without the need for retraining or parameter tuning. In other words, the key challenge is: what is an analog of natural language prompting for graph machine learning tasks? The second challenge is how to design model architecture and pretraining objectives that enable models to achieve in-context learning capability across diverse tasks and diverse graphs in the unified task representation. Existing graph pretraining methods [7, 24, 8, 13] only aim to learn good graph encoder and require fine-tuning to adapt to different tasks, while existing meta-learning methods over graphs [19, 9, 3, 17, 25] only aim to generalize across different tasks within the same graph. On the other hand, achieving in-context learning requires tackling the more difficult setting of generalizing across the graphs and tasks without finetuning.
Here we present a general approach for solving these two challenges for classification tasks on graphs: (1) prompt graph, in-context graph task representation, and (2) Pretraining Over Diverse In-Contex, Graph Systems (PRODIGY), a framework for pretraining an in-context learner over prompt graphs.
We propose prompt graph (Figure 1) to provide unified way to represent diverse node-, edge- and graph- level machine learning tasks. Promptgraph first contextualizes the input nodes/edges on which we make prediction (including both the prompt examples and the queries), then connects them with additional la- bel nodes, such that the prompt examples are interconnected with queries. Such a unified representation allows to specify diverse graph machine learning tasks to the same model regardles of the graph size.
PRODIGY then designs both model architecture and pretraining objectives with the prompt graph in-context task formulation, such that the model is pretrained to solve tasks across a wide range of tasks and graphs, and can continue do so out-of-the-box. We design a graph architecture that utilizes graph neural networks to learn node/edge representations and an attention mechanism to communicate over prompt graph. Furthermore, we propose a family of in-context pretraining objectives over prompt graph. In particular, this includes a novel self-supervised pretraining task, neighbor matching, where we classify which neighborhood a node or edge belongs to.
We use PRODIGY framework to pretrain on citation networks (MAG240M (5)) and knowledge graphs (Wiki [22]). We then show that such model (without any retraining) provides strong performance on in-context paper category classification and knowledge graph completion tasks on novel graphs it was never trained on (arXiv, ConceptNet, FB15K-237, NELL) [6, 16, 23]. Specifically, PRODIGY improves upon contrastive pretraining baselines with hard-coded adaptation for in-context setup by 18% on average across all datasets and numbers of labels to classify among. Moreover, it also outperforms standard finetuning with limited data by 32.6% on average with in in-context learning. It even outperforms the state-of-the-art few-shot learning methods trained on the testing downstream graph with pure in-context learning. Finally, we further demonstrate that our methods achieve increasingly higher performance with more examples in the prompt even beyond what was pretrained with, which shows that the model really learns to learn from context.

---

2 In-context Learning over Graphs
In this work, we specifically focus on in-context learning for node and edge classification tasks on graphs with few-shot prompting, which are the forms of the most standard and important graph machine learning tasks. this section, we introduce the concrete classification tasks over graphs and few-shot prompting over them with our in-context task representation prompt graph.
2.1
Classification Tasks over Graphs
We define graph as G= EV),R) where V, E, R represent the set of nodes, edges and relations. An edge e = (u,r,0) E& consists of subject E V, relation rE R and an object E V.
Given a set of classes V standard classification task is predicting the labeling E y of each input IE X' node- classification task is similar but each input is a single node in additional auxiliary information of the entire graph G. For example, over a citation network consisting of authors and papers, node-level classification task could be predicting the primary institution of each author. Similarly, an edge-level classification task is predicting the best labeling of potential edges formed by any pair of nodes, i.e,, X'=V V v/.. A common special case is that the classes the same as the relations y = -R, such as predicting the relation between entities over knowledge graphs. More generally, the same definitions can be extended to subgraph and graph-level classification tasks, where the input data x may consist of more nodes and edges, and essentially represents a subgraph of G.
X' =-=; V, with the
Since we are interested in tasks of different types/levels, we design unified formulation, where the space of the input X' consists of graphs. i.e., 24 E 2x,,, = (V,,E,RR,) R,). For node classification task, Gj only consists of the input node that aim to make predictions on, i.e., [V,]= 1 and [E, =0; for edge classification task, consists of (subject, object) pair, i.e., [V. = and [E =0.
2.2 Few-shot Prompting
Here we define in-context learning setup for classification tasks over graphs with few-shot prompting. For k-shot prompt with a downstream m--wa classification tasks with | = m classes, we use small number of input-label pairs S ((x;yy) } as prompt examples of the task specification, such that there are & input-label pairs with label y for each e.) We also give the model set of queries Q {x; }
that we want to predict labels for.
We emphasize an important difference of classification tasks graphs from language and other modalities. Namely, since all input datapoints are nodes/edgees/sbbgrahhs from the larger source graph G, this graph contains critical information and provides contexts for the inputs, e.g., the local neighborhood of the input node that we aim to predict. Hence, besides S and Q, we also need to include the source graph G in the prompt.
Given the above information as the prompt, the pretrained model should be able to directly output the predicted labels for each datapoint in Q via in-context learning. Thus, how to formulate the information as a unified and efficient form of input poses a unique challenge and affects the model architecture. Below, we present our in-context task formulation prompt graph designed to do so.
2.3 Prompt Graph Representation
Inspired by [2], we propose prompt graph as a unified representation of a k-shot prompt over graphs for an m-way classification task (Figure 1). A prompt graph is composed of data graphs and task graph:
Data graph. To construct prompt graph, we first perform contextualization of each datapoint ar = EV.,E,, R.) in S and Q in the source graph G to form data graphs. The goal is to gather more information about the v; from the source graph G without having to represent the entire source graph explicitly. There many potential designs for contextualization, from explicitly retrieving subgraphs to implicitly using embedding-baas based methods. Here we construct data graph G by sampling k-hop neighborhood of V; in G. In other words, G = (VP R.) ~ D Ne Wheighhoo(()) V,,, V.,9. G,1), where Vi C C CV, & C ε C E, R, C R C R, and Nei Neighbor is a function that returns the exact i-hop neighbors of each node in V/. With this data graph GP, we call the node set that corresponds to the nodes in V; before contextualization input node set, e.g., the target node to classify in node classification task and the pair of nodes in link prediction task.
i
3

---

Task graph. After contextualizing each datapoint to data graph Go, we then construct task graph G" to better capture the connection and relationship among the inputs and the labels. For each data graph Gp from the previous stage, we have data node Uz, that represents each input; for each label, we have a label node Vy.- So overall, task graph contains m m-k++n data nodes (m-k prompt examples and n queries) and m label nodes, as shown in Figure 1.
Now we add edges between the data nodes and the label nodes: know the labels of each graph, we add single directional edges from all label nodes to each datapoint in the query set, i.e., each query data node Us; will be connected to all the label nodes as shown by the yellow edges in Figure I; For the prompt examples, we connect each data node to all the label nodes, where the edge with the true labels is marked as T while the others are marked as F, as shown by the green and red edges in Figure 1 respectively.
For the query set, since do not
Together we propose the prompt graph that consists of both data graphs and task graph. Prompt graph effectively captures the re relationship between input data 2; and the label y; through the context captured in data graph GPP and the data node Un, and the label node Uy. in the task graph GT. is also possible to extend prompt graph to non-classification tasks and free-form text prompting. For example, for numerical regression (e.g. molecular energy prediction) and other free-form generation tasks (e.g. text generation), one can extend our task graph to contain vector values on the edges to represent yi- Then different label nodes would represent different prediction tasks. To support more general forms of prompting, one can include additional task information and instructions in the feature of label nodes, and additional description paired with each datapoint in the global feature in data graph.
Pretraining to Enable In-context Learning
So far given few-shot prompt for a classification task over graphs, have defined prompt graph representation for it that captures relationships between the prompt examples, queries, and labels. Now we need to design pretraining strategy that can pretrain generalizable model capable of in-context learning. We assume access to pretraining graph Spretrain that is independent of the source graph for the downstream task.
In this section, we introduce PRODIGY, a general pretraining framework over Spretzain that is designed specifically for enabling in-context learning over downstream classification tasks without any additional finetuning steps on arbitrary graphs Our framework PRODIGY has two main components: model architecture over prompt graph and in-context pretraining objectives.
3.1 Message Passing Architecture over prompt graph
Next we introduce our model architecture over the prompt graph consisting of two submodules:
Data graph Message Passing. First, we apply message passing GNN module MD that learns node representation E for nodes in each go.
E EERIViXd MG))
where d is the embedding dimension. Mp can be implemented in multiple ways, such using Graph Convolutional Network (GCN) or Graph Attention Networks (GAT) [11, 18).
To read out a single embedding G; for each data graph, we perform another aggregation step to pool node embeddings. For node classification tasks, we take the updated node representation of the single input node that we aim to predict, i.e.:
Fy=
For link prediction tasks, we concatenate the pair of nodes, which we want to predict a link between, as well as max pooling over all node representations following [10] with an additional linear projection layer at the end to convert the embedding size back to d.
G - TT
where || represents concatenation, W E Rddxd is learnable weight matrix and b is the learnable bias.
Task graph Message Passing. Note in the previous step there is no communication between different datapoints in S and Q. Now we would like to communicate between them via message passing over

---

the task graph G². We apply another GNN Mr on the task graph to obtain updated representation of data nodes and label nodes.
H=M
where H is the obtained embedding per node. The initial embedding of data node 1/2 is G, and the embedding of label node vy. can either be initialized with random Gaussian or additional information available about the labels. Each edge also has two binary features Eij that indicate 1) whether the edge comes from an example or a query, and 2) the edge type of T or F. For Mr. we use an attention- based GNN, where each node performs attention to other nodes at each layer. See the architecture detail in the appendix C.
The goal of this step is to learn better representation of the label nodes using the support examples and propagate label information back to the support and query graph representation for a more task-specific graph representation.
Prediction Read Out. Finally, we readout the classification logits O by taking cosine similarity between each pair of query graph representation and label representation, in contrastive learning:
O == paneiinessmmllarityHHHHe
(5)
Note that we could perform the two message passing steps for multiple rounds to have more communi- cation between 2; and learn better representation. One key insight is that different in-context prompt examples share information through the label nodes, which can be seen as an information bottleneck.
3.2 In-context Pretraining Objectives
In order to pretrain the model for solving the downstream graph tasks in-context, we propose a set of in-context pretraining objectives. The goal is to pretrain the graph model using a large pretraining graph Spreerain independent of the downstream task graph, such that the model can directly be applied on downstream tasks with in- mn-oonext learning.
Our main design principle is that we formulate each pretraining objective in in-context learning way. Most previous graph pretraining objectives only pretrain a shared graph encoder to perform various tasks with task-specific heads, they require finetuning for another task-specific head over each downstream task. In contrast, we explicitly construct in-context pretraining tasks in prompt graph form and pretrain the model solve diverse tasks in-context with the same set of weights, such that it can perform in-context learning directly over downstream tasks.
Below, detail our proposed family of in-context pretraining objectives in terms of three components: 1) pretraining task generation, including few-shot prompt (i.e. Figure 1(A)) and corresponding labels, 2) converting generated few-shot prompt prompt graph format (i.e. Figure 1(B,C)) with augmentation, and 3) pretraining loss over the generated prompt
grompt graph.
3.2.1 Pretraining Task Generation
We propose two methods to generate pretraining tasks from the pretraining graph Spretrain in the form of few-shot prompts: neighbor matching and multi-task.
Neighbor Matching. Given the pretraining graph, we construct self-supervised in-context pretraining tasks with the goal of classifying which local neighborhood node belongs to, where each local neighborhood is defined by the example nodes belonging to that neighborhood. Intuitively, we sample multiple subgraphs from the pretraining graph Gpretrein as the local neighborhoods, and we say a node belongs a local neighborhood if it is in the sampled subgraph.
Formally, we denote NMg. as a sampler that generates m- neighbor matching tasks, where each includes k-shot prompt (Spretrain S SNM, QNM) (see subsection 2.2 and Figure 1(A)) and the labels of the queries. For simplicity of the notation, we will include the labels in QNM as paired with the inputs:
(Spretrain:SN QNM) ~ NM (Gpretrein))
(6)
To generate these, we first sample m nodes from the pretraining graph Gpretrains where each of the sampled node corresponds to one class.
C= (c) c: Unétorm(Vpeeezzaa)
(7)
5

---

For each sampled node/class Ct we sample & different nodes from its exact (-hop neighbors, These k nodes serve as examples of label Ci. We also sample additional to form the query set. Formally,
a nodes similarly for each label
m
N; = Nei.gbehhorree (ctbboor(cc.ppeeeee (ca.bboocccc.rreriaaal
(8)
S (1.,))
2j Uniform((N)
(9)
2 (<)),)
Ij~ Uniform
(10)
In such a way, constructed neighbor matching pretraining task sample in the format of a few-shot prompt U
The neighbor matching task generation process outlined above is specifically applicable when the downstream tasks are also node classification. When the downstream task link prediction, we may adapt the above neighbor matching tasks to over edges correspondingly. Specifically, each sampled input node z; to an edge by randomly sampling an edge that contains ø;. Then, instead of classifying which neighborhood a node in the query set belongs, now the neighbor matching task is to classify to which neighborhood an edge in the query set belongs.
expand
Multi-task. When the pretraining graphs have node or edge-level labeling f(x::)= E ) for some Jy Ppretrsin or Epretrain: we can further leverage this signal to perform supervised pretraining. Similar to neighbor matching, the key is to construct such supervised pretraining tasks in the format of shot prompts and corresponding labels.
(Gpretrain SMr SMI,SM)) (Trr)~MT,,. M (9pretrain:)
(11)
For node classification tasks, we first sample m labels from the whole label set. set. Then, for each label, we directly sample & nodes support examples and m nodes with labels in this set as query examples.
C {c S, = (()),),=))} Q = -,=((,,,,,=,),
~~ Uniform())
(12)
Iy~ Uniform(((x:](((=====0
(13)
Zyr~ Uniforn ((x
(14)
US, We then construct a task with the few-shot prompt as Sppretrein:SMT QMT UQ)) For link prediction, we directly use the edge type function as f. f.c.c f((11,02))=r =r + (11,5,02) E With this f, we may directly sample m edge types and construct pretraining tasks in similar way as above.
The benefit of such a supervised pretraining objective is that it could directly resemble the format of downstream tasks, compared with neighbor matching objective, which may only serve surrogate. However, it requires extra labels if f is not part of Gpretrain. e.g. node classification labels that may not exist for some pretraining graphs.
3.2.2 Prompt graph generation with augmentation
After we obtained the few-shot prompts and labels for either of the two tasks (NeighborMatthing and multi-task), we need to construct the prompt graph for pretraining. addition to the standard construction process described in subsection 2.3, we add additional augmentation step to augment
g
Prompt Examples 5: Input Label y
Queries Q:
hops liihhoors
G2
Figure 2: Neighbor matching pretraining task generation. The key idea neighbor matching is to classify whether node/edge is in the local neighborhood of set of sampled nodes (as labels). Given a set of sampled nodes and their two-hop neighbors as prompt examples, we aim to classify whether query node is also two-hop neighbor of each sampled node.

---

the data graphs as inspired by Contrastive Learning. The key insight is to corrupt data graph such that the pretrained model learns representation invariant to various corruptions.
Here we demonstrate how we adopt graph augmentation techniques during the construction of prompt graph from a few-shot prompt generated from Sprotrain. We first still sample the k-hop neighhor sub-
graph of each sample G; in the prompt examples and queries: CP ~ ① j= (S.
Then we adopt the following two augmentation techniques to create augmented data graph g, including (1) node dropping, and (2) node feature masking [24]. For node dropping. we randomly drop nodes from the k-hop neighbor subgraph and take the remaining graph as G -DropNode (CP)) For node feature masking, we randomly mask the feature of subset of nodes with value zero to create g,a0 MaskNode PP). With the augmented data graphs for each datapoint in the prompt examples and the queries, we may accordingly construct the task graph G" by creating a data node Ua, for each augmented data graphs and the label node Vy, introduced in subsection 2.3. Combining data graphs with task graph, we obtain the prompt graph formulation with augmentation for the few-shot prompt.
3.2.3 Pretraining Loss
Finally, we pretrain the model with the cross-entropy objectives over generated prompt graphs:
(Gpratrain,SNM,
NM (Sprotras))
(15)
(Spretrain SMT, ~ MT (Spretrain:
(16)
C E
CE(ONM...NNM.) + CE(OMT,,::NT..)
EE
(17)
2, QMI
where O is the logits produced by our model over input of 9; and G" produced from 2NM, as described in subsection 3:1; 1/M, is the corresponding label of x;; in QNM; Similar for MT terms.
4 Experiments
4.1 Experimental Setup
Datasets.
For pretraining, we use two datasets: MAG240M [5], a large scale citation network with 122 million nodes and 1.3 billion edges; and Wiki, a knowledge graph (KG) constructed from Wikipedia [22] with 4.8 million nodes and 5.9 million edges. After the model is pretrained we evaluate its in-context learning capability on diverse classification tasks over 4 graphs: arXiv [6]. ConceptNet [16], PB15K-233 [23], NELL [23]. We use subsets of knowledge graph datasets same as in [10, 23]. For arXiv, the downstream task is an m-ways node classification task that predicts the paper category. For knowledge graph datasets (ConceptNet, FB15K-231 NELL), the downstream task an m-ways relation type classification task that predicts the relationship connecting the two input nodes.
Evaluation. We pretrain our model on MAG240M and Wiki and then we evaluate the in-context learning performance on differnt downstream datasets that belong to similar domain as the pretraining dataset (e.g., pretraining on Wiki and evaluating on ConceptNet, FB15K-233 and NELL). Fach of the downstream classification datasets has its original train, validation, and test splits. To simulate the situation where there are limited amount of labeled data in the downstream task, we randomly select 10 nodes (or edges) from the training split per way as the prompt examples with known labels. Then, we construct a k-shot prompt for test nodes (or edges) from the test split by randomly selecting k examples per way from these available examples. This allows us to test the model's ability to learn in-context relationships and perform well on classification tasks with truly limited known labels. By default we use k= 3 shots in our experiments.
Methods and Baselines. We consider three versions of our proposed framework PRODIGY: I)
PG-NM, which uses neighbor matching task for pretraining: 2) PG-MT, which employs multi-task pretraining; and 3) full PRODIGY, which combines the previous two methods. To augment the data, use DropNode and MaskNode augmentations with probability 0.5 per node for each method.
We consider three baselines for comparison: 1) NoPretrain, which uses a randomly-initializeed model with the same architecture our pretrained models; 2) Contrastive [24], which employs a standard contrastive learning method with the same augmentation as above and uses a hard-coded nearest neighbor algorithm to adapt to our in-context learning setting. Specifically, we classify the query by
7

---

Table I: In-context learning accuracy (%) on arXiv paper category classification on 500 sampled test tasks with 3-shot prompts. PRODIGY was pretrained on MAG240M and is then applied in-context to arv,v, which has completely different structure and a different set of paper categories. PG-NM and PG-MT are ablations of PRODIGY.
Classes
NoPretrain
Contrastive
PG-NM
PG-MT
PRODIGY
Finetune
3
33.16 16 20:30
65.08 2034
72.50 2
65.64 2033
73.09 = 0.36
65.42 5.533
5
18.33 2021
51.63 2029
61.21 +
51.97 2:27
61.52 0.28
53.49 461
10
9.19 20.11
36.78 2.19
46.12 0.19
37.23
0.20
46.74 = 0.20
30.22 3.77
20
4.72 0.06
25.18
2.11
33.71 0.12
25.91
0.12
34.41
17.68 4 1.15
40
2.62 0.02
17.02 0007
23.69 0.006
17.19 1: 0.08
25.13 =007
8.04 = 3,00
Table 2: In-context learning accuracy (%) on ConceptNet, FB15K-237 and NELL (from top to bottom) on 500 sampled test tasks with 3-shot prompts. PRODIGY was pretrained on Wiki, which has completely different node and relation types from graphs it is then applied on in-context.
Classes
NoPretrain
Contrastive
PG-NM
PG-MT
PRODIGY
Finetune
4
30.4 1063
44.01
1.1
46.94 161
51.78
0.63
53.97 0.6
53.85 9.29
5
33.54 1061
81.35 + 08
80.35 0.57
89.15
0.46
88.02
0.48
82.01 12.83
10
20.0 1:35
70.88 =: 0.48
71.68 0.45
82.26
0.40
81.1 0.39
71.97 = 6,16
20
9.2 + 0.18
59.8 + 0.35
59.9 + 0.35
73.47 0.32
72.04 0.33
64.01 +: 4.66
40
2.5 + 0,08
49.39 + 0.23
46.82 + 0.1
58.34
0.22
59.58 0.22
57.27 +: 3,33
5
20.95 0.52
83.38 0.
82.39
85.26
88.09
87.22 +
10
11.0 0:26
74.54 +: 0.46
75. 14
:: 0.43
78.15 0.41
82.47 +: 0.39
71.90 + 5.90
20
5.34 0.13
65.68 034
65.68 0.34
68.38
74.72 0.31
66.19 + 8.46
40
2.5 +: 0.06
56.7 0.23
54.91
0.22
51.24
1: 0.25
60.04 0.23
55.06 + 4.19
comparing its pretrained embedding against the average embedding of the example inputs of each class. 3) Finetune [7], which trains an additional linear classification head on top of the graph encoder pretrained with contrastive learning, following the standard practice.
4.2 In-Context Learning Results
We first evaluate the in-context learning capability for node classification and link prediction with various numbers of ways (i.e. number of classes to classify among). The results are presented in Table 1 and for node classification and link prediction.
Strong in-context learning performance. The results demonstrate that our method PRODIGY consistently outperforms all other baselines in this setting. It achieves the highest average accuracy across all ways on arXiv, with an average improvement of 28.6% and up 48% over the best baseline of Contrastive. Over KGs, PRODIGY also outperforms contrastive learning on average by 12.2%. PRODIGY also demonstrates smillarto-better performance compared to Finetune, which requires additional training on downstream tasks. On arXiv, we see an average improvement of 77.7% over all ways. This can be attributed to the diverse set of pretraining tasks incorporated PRODIGY, which allows the model to avoid overfitting on specific tasks and learn in-context.
Self-supervised pretraining PG-NM bridges different tasks. In particular, we highlight that the pure sulf-super pretraining method PG-NM produces significantly higher in-context learning performance over arXiv than baselines, even though the model is pretrained different tasks from the downstream task. This advantage can be further leveraged by pretraining on even larger-scale unlabeled datasets. On the other hand, PG-MT follows the supervised pretraining objective that directly resembles the format of downstream tasks. On KGs, this allows PG-MT to adapt better to downstream task even sometimes compared the full PRODIGY marked by underlines), while PG-NM might have overfitted to the incorrect strategy of only identifying co-occurring nodes. Yet, PG-MT performs worse on arXiv potentially due to less diversity. The full PRODIGY, which ensembles the two, achieves more diversity than either single task and therefore achieves the best performance over both worlds.
Outperforming meta-learning method trained on test graph. Finally, we compare PG-NM in-context learning performance against state-of-the-art meta-learning method TENT [20] over the downstream test graph ar Xrvi. We evaluate the average 3-ways classification tasks performance over only test labels, since TENT trains on train labels from arXiv. PG-NM achieves 69.07 99.07% over the 65. 13% of TENT, even though PG-NM has never been trained any paper category classification
8

---

0.64
0.60
accuraty
0.56
accuracy
0.55
0.52
0.48
0.44
Contrastive PRODIGY
Contrastive
0.36
PRODICY
8
12
16
20
24
28
2500 5000 7500
10000 12500
17500 20000
training steps
Figure 3: In-context learning accuracy on Concept- Net in a 4-ways setting wrt. the number of prompt examples (shots).
Figure 4: With more training steps and data arXiv (5-ways), PRODIGY keeps improving while the baseline (Contrastive) saturates.
task during pretraining, This demonstrates the power of self-supervised pretraining over large amount of data compared to supervised meta-learning over the limited labeled data (train labels in arXiv).
4.3 Ablations
Aside from PG-NM and PG-MT, we also conduct ablation studies on various configurations of the self- supervised objective PG-NM as described in 3.2. See the full results in Appendix E and Table 4. Overall, the ablation results reveal that using of the elements together results the highest performance. Specifically, attribute prediction (see appendix A) has the greatest impact on PG-NM's performance, as its removal results in an average 7% drop across ways, shown in the 'No-Attr' column.
4.4 Evaluation using different numbers of in-context examples
We investigate our method's s ability to learn from the context by analyzing its performance as the number of prompt examples changes. Figure shows the result on ConceptNet. See full results on other datasets in Appendix F. As the number of prompt examples increases, the margin of our proposed PG models over the baseline increases. This supports the hypothesis that the PRODIGY models can more effectively learn the unknown task by reasoning about the common characteristics of prompt examples.
4.5 Scaling with Data Size
Finally, explore how the model scales with more pretraining data. The result on arXiv in 5-ways set- ting is illustrated in Figure 4. It shows that the Contrastive baseline saturates quickly and its performance fluctuates as trained over more pretraining data. Instead, PRODIGY consistently shows an improvement in performance as more data is pretrained on, since the pretraining tasks are harder and more diverse.
5
Related Work
5.1 In-context Learning of Large Language Models
Pretrained large language models can make predictions for diverse downstream tasks directly by prompting with a few examples of the task or more generally any textual instructions. This ability called in-context learning. Comparing to previous language encoder models like BERT [4], drasti- cally reduces the adaptation effort comparing to fine-tuning, and has demonstrated strong performance broad range of models and tasks. Our work extends this success similarly to graph data compared to the current pretrained graph encoders, such that a single pretrained model can be adapted to different classification tasks over different graphs without additional fine-tuning but only few-shot prompting.
5.2 Pretraining on Graphs
There are many existing works on pretraining over graphs[7, 24, 8, 13]. However, they all follow the general paradigm of learning good graph encoder that can perform certain pretraining tasks, such as masked feature prediction [7] and paired graph classification [24]. To adapt to any downstream tasks, then requires finetuning classification head on top of the encoder with large amount of task specific data for each downstream task. In contrast, we explore pretraining methods for inducing general in co-tent learning ability, such that the pretrained model can be directly used for various downstream tasks with no gradient updates.

---

5.3 Meta Learning on Graphs
Another closely related line of works is meta-learning methods over graphs that aim to address standard few shot learning problems over graphs[19, (1,, 9, 25J. However, existing meta-learning methods are only designed and tested for generalizing across different tasks on the same graph: the methods are trained set of training tasks on graph, then tested over disjoint but similar set of test tasks over the same graph. They are shown to exhibit optimal performance only when trained on similar curated tasks [10]. Different from this, our work explicitly focuses on the in-context learning performance, i.e. model performance on graphs and tasks completely different from the pretraining without additional fine-tuning.
Conclusion
We introduce PRODIGY, the first framework that enables in-context learning on graphs. A model that is pretrained using PRODIGY can seamlessly execute new classification task over new graphs represented by prompt graph. It markedly surpasses the performance of other baseline models with in-context learning, even those that employ finetuning, in both the node and edge classification tasks.
Acknowledgments and Disclosure of Funding
We thank Camilo Ruiz and Michael Moor for discussions and for providing feedback on manuscript.
We also gratefully acknowledge the support of DARPA under Nos. HR001 121 2900039 (TAMI), N660011924033 (MCS); ARO under Nos. W91 16-1-4221 (MURI, W91 INN-IN------ 16-10161 (DURIP); NSF under Nos. OAC-1835598 (CINES), OAC-1934578 (HDR), CCF-1918940 (Expeditions), NIH under No. 3U54HG010426-044
(HuBMAP), Stanford Data Science Initiative, Wu' Neurosciences Institute, Amazon, Docomo, GSK, Hitachi, Intel, JPMorgan Chase, Juniper Networks, KDDI, NEC, and Toshiba. Hongyu Ren is supported by Masason Foundation PhD fellowship, Apple PhD fellowship and Baidu PhD scholarship. Qian Huang supported by Open Philanthropy Al fellowship.
The content solely the responsibility the authors and does not necessarily represent the official views of the funding entities.
10

---

References
[1] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[2] K. Cao, J. You, J.d Leskovec. Relational multi-task learning: Modeling relations between data and tasks. In International Conference on Representation Learning (ICIRR, 2022.
[3] M. Chen, W. Zhang, W. Zhang, Q. Chen, and H. Chen. Meta relational learning for few-shot link prediction in knowledge graphs. arXiv preprint arXiv:1909.0155: 2019.
[4] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. ArXiv, 8000000,0 2019.
[5] W. Hu, M. Fey, H. Ren, M. Nakata, Y. Dong, and J. Leskovec. Ogb-Isc: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2103.094300 2021.
[6] W. Hu, M. Fey, M. Zitnik, Y. Dong, H. Ren, B. Liu, M. Catasta, and J. Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint trxiv:2005.00687 2020.
[7] W. Hu, B. Liu. J. Gomes, M. Zitnik, P. Liang, V. Pande, and J. Leskovec. Strategies for pre-train pre-training graph neural networks. arXiv preprint arXiv: 12265,, 2019.
[8] Z. Hu, Y. Dong, K. Wang, K.-W. Chang, and Y. Sun. Gpt-gnn: Generative pre-training of graph neural networks. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 1857-1867, 2020.
[9] K. Huang and M. Zitnik. Graph meta learning via local subgraphs, 2020.
[10] Q. Huang, H. Ren, and Leskovec. Few-shot relational reasoning via connection subgraph pretraining. Advances in neural information processing systems, abs/2210.06722, 2022.
[11] T. N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks, 2016.
[12] Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemooyrr and V. Stoy- anov. Roberta: A robustly optimized BERT pretraining approach. CoRR, 11662 2019.
[13] J. Qiu, Q. Chen, Y. Dong. J. Zhang, H. Yang. M. Ding, K. Wang, and J. Tang. Gec: Graph con- trastive coding for graph neural network pre-training. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pages 1150-1160,200 2020.
[14] S. Sharma and R. Sharma. Identifying possible rumor spreaders on twitter: A weak supervised learning approach. In 2021 International Joint Conference on Neural Networks (IJCNN), pages 1 1-8, 2021.
[15] K. Song, X. Tan, Т. Qin, J. Lu, and T.-Y. Liu. MPNet: Masked and Permuted Pre-training for Language Understanding. In Advances in Neural Information Processing Systems, volume 33, pages 16857-16867. Curran Associates, Inc., 2020.
[16] R. Speer, J. Chin, and C. Havasi. Conceptnet 5.5:
An open multilingual graph of general knowledge. In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017.
[17] J. Sun, Y. Zhou, and C. Zong. One-shot relation learning for knowledge graphs via neighborhood aggregation and paths encoding. Transactions on Asian and Low- Resource Language Information Processing, 21(3):1-19, 2021.
[18] P. Veličković, G. Cucurull, A. Casanova, A. Romero, P. Lid, and Y. Bengio. Graph attention networks, 2017.
[19] S. Wang, K. Ding, C. Zhang, C. Chen, and J. Li. lask-adaptive few-shot node classification. 2022.
[20] Wang, K. Ding, C. Zhang, C. Chen, and J. Li. Task-adaptive few-shot node classification. Pro- ceedings the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, 2022.

---

[21] S. Wu, W. Zhang, F. Sun, and B. Cui. Graph neural networks in recommender systems: A survey, ACM Computing Surveys, 55:1 37, 2020.
[22] W. Xiong, M. Yu, S. Chang, X. Guo, and W. Y. Wang. One-shot relational learning for knowledge graphs, 2018.
[23] W. Xiong. M. Yu, S. Chang, X. Guo, and W. Y. Wang. One-shot relational learning for knowledge graphs. In EMNLP, 2018.
[24] Y. You, T. Chen, Y. Sui, T. Chen, Z. Wang, and Y. Shen. Graph contrastive learning with augmentations. Advances in neural information processing systems, 33:5812-5823, 2020.
[25] C. Zhang, H. Yao, C. Huang, M. Jiang, Z. Li, and V. Chawla. Few-shot knowledge graph completion. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 3041-3048, 2020.
12

---

A Attribute Prediction Loss
For each augmented DataGraph gave, the certain node features F, are masked during MaskNode augmentation. Therefore, we can reconstruct them using the learned embedding E, with a MLP and train with MSE reconstruction Loss.
1 Σ MSE(F.,,LP(E.). 100
Lattr( (goug)
B Dataset Statistics
Table 3: Dataset statistics
Dataset
Nodes
# Edges
# Classes
MAG240M
122M
1.3B
153
Wiki
4.8M
5.9M
639
arXiv
169K
1.2M
40
ConceptNet
791K
2.5M
14
FB 15K-237
15K
268K
200
NELI
69K
181K
291
C Task Graph GNN Architecture
For the GNN over task graph Mr, attention- based GNN, where each node performs attention to other nodes at each layer:
Bu - MLP=MT H! Hl)les)
(18)
exp(B1g)
(19)
H ReLU
BN
HW Σ
aij W H
(20)
jEN(jU(i)
D Hyperparameters
D.1 Model Architecture, MAG240M and arxiv
We initialize the node features in citation network datasets using pretrained language model (RoBERTa [12] base model trained on NLI and STSB).The architecture of our PromptGraph model in all of our proposed methods for citation network datasets (full PRODIGY, PG-NM. and PG-MT) and the baseline (NoPretrain). consists two message passing layers, Mp, over the DataGraph and one message passing layer, MM, over the TaskGraph. These layers defined in Section 3.1.
For the Contrastive method, the architecture includes two message passing layers, Mas over DataGraph, and contrastive learning component that is defined in Section 4.1. The mode for Finetune the same as the Contrastive method, with the addition of linear layer head over the output of the two Mg layers, also described in Section 4.1.
D.2 Model Architecture, knowledge graph datasets
We initialize node and edge features knowledge graph datasets using pretrained language model (MPNet [15]). The architecture of our PromptGraph model in all of our proposed methods for knowledge graph datasets (full PRODIGY, PG-NM, and PG-MT) and the baseline (NoPretrain), consists of two message passing layers, Mp. over the DataGraph, aggregator described by
13

---

Table 4: Ablation of PG-NM on arXiv.
Ways
PG-NM
3 - shot
No Attr
No Aug
No Attr, Aug
No Attr, Aug, Mr
3
72.50 0.35
69.13 + 1.09
65.74 + 1.12
68.98 + 108
66.53 + L12
63.60 = 1.06
5
61.21 + 0.29
57.49 + 0.92
52.78
+ 0.90
57.50 + 0as
53.89 + 0.2
51.27 + 0.89
10
46.12 0.19
42.03 + 0.60
37.99
+ aua
42.43
064
38.87 + 0.59
37.62 + 0.34
20
33.71
+ 0.11
30.18
+ 0:38
26,60 + 0.36
30.89
+ 0.38
27.50
+ 0.36
27.44 + 0.17
40
23.69 0.07
21.44 + 0.22
18.03 +
21.97 o.44
18.52 + 0.22
19.69 + 0.8
Equation 3, and two message passing layers, Mr, over the TaskGraph, which only pass messages along the positive and query edges. These layers are defined Section 3.1.
For the Contrastive method, the architecture includes two message passing layers, Ma. over the DataGraph, an aggregator described by Equation 3 and contrastive learning component that is defined in Section 4.1. The mode for Finetune is the same as the Contrastive method, with the addition of linear layer head over the output of the two Ma layers, also described in Section4.1.
D.3 Training, MAG240M
The following describes our pretraining setup:
The pretraining task we used consisted of 30 ways, 3 shots, and 4 queries per task. This specific task configuration was carefully selected strike a balance between complexity and diversity the training data, without overwhelming the GPU memory.
We checkpoint the model every 500 steps.
Our pretraining setup included a model with an input dimension of 768 and an embedding dimension of 256, batch size of 1, and the AdamW optimizer with a learning rate of x 10-3 and weight decay of 1 x 10 a pretraining task with 30 ways. shots, and queries per task, and checkpointing every 500 steps. This consistent configuration was applied across all the methods for fair comparison. Our full PRODIGY setup, on average, involves sampling Neighbor Matching task per I multitask pretraining tasks.
For our evaluation process, we computed zero-shot transfer performance of the model on the test set, using the pretraining checkpoint at the 10,000 step pretraining. The evaluation was conducted on 500 test tasks, with batch size of 5, measured on the downstream task of graph classification accuracy. To maintain consistency, we kept the number of shots and queries constant at for all evaluation tasks.
D.4 Training, Wiki
The following describes our pretraining setup:
Our pretraining setup included model with an input dimension of 768 and an embedding dimension of 256, the AdamW W optimizer with a learning rate of x 10-3 33 and weight decay of x 10-3 pretraining task with 30 ways, 3 shots, and 4 queries per task, using batch size of 10. and checkpointing every 500 steps. This specific task configuration was carefully selected to strike a balance between complexity and diversity in the training data, without overwhelming the GPU memory. This consistent configuration was applied across all the methods for fair comparison. Our full PRODIGY setup involves sampling one neighbor matching task per 50 multitask pretraining tasks.
For our evaluation process, we computed zero-shot transfer performance of the model on the test set, us- ing the pretraining checkpoint at the 8,000 step of pretraining. The evaluation was conducted on 500 test tasks, with batch size of 1, measured on the downstream task of graph classification accuracy. To main- tain consistency, we kept the number of shots and queries constant at 3 for all evaluation tasks. We sam- ple h-hop neighbourhoods for ConceptNet, FB15K-237 and NELL and 2-hop neigbourhoods for Wiki.
E Ablation on Table 4 for the PG-NM setting
In Table 4, we ablate on various configurations of the self-supervised objective PG-NM. As described in Section 3.2, PG-NM composed of an attribute prediction loss, dropnode and zeronode augmentations, with the default setting of sampling shots neighbor matching tasks. Our best setting, referred to
14

---

as simply "PG-NM", is also shown in Table I and comprises of attribute prediction, dropnode and zeronode augmentations, with the default setting of 3 shots.
The ablation results reveal that using all of these elements together results in the highest performance. Specifically, attribute prediction has the greatest impact PG-NM's performance, its removal results in an average 7% drop across all ways, as shown in the 'No-Attr' column.
Removing the dropnode and zeronode augmentations results in average 3% drop across all ways, as shown the No Aug column. Removing both attribute prediction and augmentations results in perfor- mance that is similar to just removing attribute prediction alone, which also roughly a 7% drop across all ways, as shown in the 'No Attr, Aug' column. Additionally, we found that decreasing the number of shots to 1 from the default setting of 3 resulted in average 3.5% drop across all ways, shown.
Evaluation using different numbers of shots
show evaluation using different numbers of shots, as shown in Figures 3, 5, and 6, 7.
FB15K-237 20 way
0.72
accuracy
0.68
0.64
0.60
0.56
Contrastive PRODIGY
0.52
8
12
16
20
24
28
shots
Figure 5: In-context leamnngg accuracy on FB15K-237 a 20-ways setting wrt. the number of prompt examples (shots).
NELL way
0.84
0.80
0.76
accuracy
0.72
0.68
0.64
0.60
Contrastive
0.56
PRODIGY
12
16 shots
20
Figure 6: In-context learning accuracy NELL in 20 20-ways setting wrt. the number of prompt
examples (shots).
G Scaling with Data Size
explore how the model scales with more pretraining tasks. Note that we use the number of train steps as proxy because the model sees more pretraining tasks as the training proceeds with almost no redundancy 20% for 10k tasks). The result on arXiv in 5-ways setting illustrated Figure It shows that the Contrastive baseline saturates quickly and its performance fluctuates given more
15

---

arXiv 5 way
0.70
accuracy
0.60
0.55
Contrastive
PRODIGY
5
10
15
20
25
30
shots
Figure In-context learning accuracy on arXiv in a 5-ways setting wrt. the number prompt examples (shots).
pretraining tasks. Instead, PRODIGY consistently shows an improvement in performance more data is pretrained on.
H
Compute
We use one NVIDIA A100-SXM4-80GB GPU for all our experiments. One pretrain run 10k steps takes 3 to 4 hours.
I Broader Impacts
Our work aims to extend the success of in-context learning to graphs and start building toward graph foundation models. This would allow cost-effective and accurate predictions, especially in domains where labeled data is scarce and long tail such as network anomaly detection, rare disease diagno- sis/treatment, supply chain disruption, and recommendations for new users. However, overreliance on prior knowledge from pretraining could also lead to increased social bias and unfair benefits to the dominate groups. To mitigate this, pretraining data should be diverse and well-balanced, and the pretrained models should be tested on downstream tasks over different groups and subdistributions.
16


Reviews:
[{"review_id": "GZEnaosPUu", "reviewer": "Reviewer_GCaE", "summary": "This paper proposed a framework for graph in-context learning. The PRODIGY architecture consists of the prompt graph, task graph, and in-context learning pretraining objective. The PRODIGY can directly conduct downstream tasks without finetuning and shows strong performance on downstream classification tasks.", "questions": "The proposed PRODIGY requires pretraining and testing on the same type of graph, such as MAG240M for pretraining and arxiv for testing, where both datasets are citation networks. It would be intriguing to observe PRODIGY's performance when the training and testing data belong to different domains.", "limitations": "N/A", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper is overall well-organized and well-written.\n2. To the best of my knowledge, this is the first work on graph in-context learning, which can be directly applied to downstream tasks without tuning.\n3. The proposed PRODIGY shows strong performance, significantly outperforming baseline methods.", "weaknesses": "1. missing important self-supervised graph learning baselines, such as GraphMAE [1].\n2. Since there is much recent progress in graph contrastive learning, including data augmentation and architecture design, I strongly encourage the authors to select more recent and competitive baselines.\n3. The authors only report results of node classification on one dataset, i.e., arxiv dataset.\n\n[1] Hou, Z., Liu, X., Cen, Y., Dong, Y., Yang, H., Wang, C., & Tang, J. (2022, August). Graphmae: Self-supervised masked graph autoencoders. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (pp. 594-604).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and valuable feedback. Below, we clarify a number of important points raised by the reviewer.\n\n> Re: contrastive learning baselines\n\nThank you for the reference. The reviewer suggests we compare our method with more contrastive learning methods. Here we would like to emphasize that our main contribution is to propose the first pretraining framework that enables in-context learning over graphs. We are not simply designing a new contrastive learning algorithm, but rather the framework we proposed enables a given contrastive learning algorithm to learn in context for graph tasks. Hence we view contrastive learning algorithms to be orthogonal contributions from ours. We will implement more contrastive learning methods and do a thorough evaluation.\n\n> Re: training and test data from different domains\n\nThank you for the suggestion. Our experiments on link prediction have demonstrated transfer across datasets from multiple domains. Enabling pre-training and testing from different domains is very challenging and requires modeling capacity to capture the gap between both node/edge features as well as the graph structure. This is exactly our future work for a more fundamental graph foundation model that can transfer and complete tasks in-context from diverse domains and even diverse tasks (node/link/graph-level prediction tasks).\n"}, {"review_id": "DTJJGekOUe", "reviewer": "Reviewer_DFam", "summary": "I have read the author’s rebuttal, I think I misunderstood the in-context learning mentioned in this paper and see the difference from other typical ICL works. I have no object to accept the paper if AC thinks it is enough contribution. \n\nThe paper introduces an in-context few-shot prompting approach for edge classification over graphs using the PRODIGY framework. The idea is to use GNN to create few shot prompts that can make LLM do better in context learning. The pretrained model exhibits strong in-context learning performance on downstream tasks, surpassing contrastive pretraining baselines and standard finetuning methods.", "questions": "See weakness.", "limitations": "NA", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Originality: The paper demonstrates some novelty by utilizing a graph to generate few-shot prompting examples for classification.\nQuality: The results presented show significant improvements over the baseline methods.\nClarity: The methodology is described clearly, providing a clear understanding of the different components involved in task construction.\nSignificance: The paper appears to lack a truly novel contribution, instead combining multiple existing approaches and claiming improvement by combining these strategies. E.g. using graphs to construct a few shot examples. How does this differ from superICL which is a more generic model that uses different downstream models for constructing in-context examples?", "weaknesses": "The paper is not very easy to read with some improvement room for the writing and fluency. For example, there are some incomplete sentences in the writing.  \"give music product 29 recommendations on spotify when being trained on Amazon book purchasing graph.\"\n\nThe paper appears to lack a truly novel contribution, instead combining multiple existing approaches and claiming improvement by combining these strategies. E.g. using graphs to construct a few shot examples. How does this differ from superICL which is a more generic model that uses different downstream models for constructing in-context examples?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and valuable feedback. Below, we clarify a number of important points raised by the reviewer.\n\n> Re: Novelty\n\nThe reviewer claims that the paper lacks a truly novel contribution. We respectfully disagree. The paper proposed one of the first frameworks that allow for in-context learning over graph tasks. We are not using graphs to construct few-shot examples, but rather propose a way to unify the formulation of different types of graph tasks, including node classification and link prediction, such that the model is able to learn in context. \n\n> Re: how does this differ from superICL?\n\nWe would appreciate it if the reviewer can provide a detailed reference to the paper instead of only giving a name. To the best of our knowledge, we find this paper [1] by looking up the model’s name superICL. However, we fail to find connections between our work (which focuses on enabling in-context learning for graph tasks) and superICL [1] (which leverages a combination of an LLM with smaller models to perform supervised tasks efficiently for in-context learning on text tasks). Please kindly let us know how these two are related.\n\nWe will improve the writing and polish the narrative in the final version.\n\n[1] Xu, Canwen, et al. \"Small models are valuable plug-ins for large language models.\" arXiv preprint arXiv:2305.08848 (2023).\n"}, {"review_id": "8O9atX9IbO", "reviewer": "Reviewer_FSU3", "summary": "\nThis paper introduces Prodigy, a method aimed at facilitating 'in-context learning' over graphs.\n\nThe key contribution of this work is the formulation of a 'prompt' that can be utilized for in-context learning with graphs. This 'prompt' is defined as a data graph which incorporates typical (input, output) examples, akin to those used in few-shot prompting setups. Given that each node in the graph serves as an input to the method, it is crucial to contextualize each node. For this, the authors add a neighborhood to each node in the data graph. This data graph is then paired with a task graph. The task graph, designed to link different parts of the prompt (e.g., connecting nodes in the data graph that are part of the same class), comprises one node for each node in the data graph (referred to as 'data nodes'), and one node for each label (termed 'label nodes').\n\n\n\nWhen running inference on a new example, the procedure begins with a Graph Convolutional Network (GCN) style message passing over the data graph, followed by the task graph. Eventually, the label is predicted based on the similarity between the representations of the query and label nodes in the task graph.\n\nTo train this model, the authors propose two self-supervised pre training objectives: neighborhood prediction and a combination of link-prediction and node prediction.\n\nThe experimental results on citation graphs and commonsense graphs indicate the potential of the Prodigy method, with it outperforming strong baselines, including fine-tuning.\n", "questions": "Q1: *L264: Then, we construct a k-shot prompt for 264 test nodes (or edges) from the test split by randomly selecting k examples per way from these available 265 examples. This allows us to test the model’s ability to learn in-context relationships and perform well on 266 classification tasks with truly limited known labels. By default we use k= 3 shots in our experiments.*\n\nGenerally, in a few-shot scenario, a fixed number of examples are included in the prompt, with the test example added at the end and supplied to the model for inference. However, the description here, which refers to the selection of both training and test examples, seems somewhat unclear. Could you please provide more clarity?\n\n\nQ2: Considering the baselines, is it the most effective approach to train a single Multi-Layer Perceptron (MLP) on top of a graph encoder? Wouldn't training the entire graph encoder end-to-end yield more effective results? \n", "limitations": "Please see weaknesses. ", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The concept of extending few-shot learning to graphs is both intuitive and attractive. Based solely on the coherent design for graph-based few-shot learning, the paper is worth considering for acceptance.", "weaknesses": "The title's use of the term “in-context learning” is questionable. While the term is currently popular due to the rise of Large Language Models (LLMs), it may inadvertently mislead readers about the actual contributions of this paper.\n\nTraditional in-context learning:\n\n* Applies to a large number of tasks\n\n* Adapts to new domains\n\n* Allows fluid task definition\n\nHowever, none of these attributes seem to apply to the proposed model. Why call it in-context learning when it essentially seems to be K-shot link prediction and node classification?\n\nThere's a crucial factor to consider: few-shot learning enables tasks to be undertaken even when no training data is available. However, PRODIGY appears to necessitate pre-existing data to operate effectively as the experiments show. This mismatch in the implementation and title is the main reason for my somewhat low score. In general, I would strongly recommend a change in title for this work (at the minimum, changing _Enabling_ to _Towards_).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and valuable feedback. Below, we clarify a number of important points raised by the reviewer.\n\n> Re: Difference from K-shot prediction.\n\nThe reviewer raises a concern on our difference from K-shot prediction. As in the related work section, most existing few-shot learning works are designed and tested for generalizing across different tasks on the same graph. They are shown to exhibit optimal performance only when trained on similar curated tasks. Our major contribution is to relax such constraints, and enable GNNs to perform both node-level and link-level tasks across drastically different domains (citation, wikipedia, freebase, commonsense) for train/test without the need to perform additional finetuning. Overall we think Prodigy is a solid contribution to explore how we can possibly formulate and also learn in context for various graph tasks across domains. We are happy to update the title to better reflect our contribution in our final version.\n\n> Re: Generally, in a few-shot scenario, a fixed number of examples are included in the prompt, with the test example added at the end and supplied to the model for inference. However, the description here, which refers to the selection of both training and test examples, seems somewhat unclear. \n\nYes, your understanding is correct. Given a test datapoint, we sample k training datapoints for each of the m classes/ways. Together these serve as our prompt. Note there is no test datapoint in the prompt. We will rephrase the sentence in our final version to make it easier to understand.\n\n> Re: is it the most effective approach to train a single Multi-Layer Perceptron (MLP) on top of a graph encoder? Wouldn't training the entire graph encoder end-to-end yield more effective results?\n\nThe reviewer wonders if doing end-to-end training can yield better results. We have run both configurations, and we did not notice a drastic difference in performance."}, {"review_id": "URP6qkhPBQ", "reviewer": "Reviewer_GRdC", "summary": "This paper proposes a pretraining framework that enables in-context learning on graph classification tasks (maybe diverse graph machine learning tasks). Specifically, it proposes a prompt graph as a unified representation for diverse tasks, then it designs a graph neural network architecture over the prompt graph and a corresponding family of in-context pre-training objectives. The experiments show that the pretrained model using this framework exhibits good in-context learning performance over various new tasks in the same domain of pre-training without finetune.\n", "questions": "1. The paper demonstrates the effectiveness of the proposed approach using the evaluation datasets from the same domain with the pre-training dataset. Does the success of this approach hinge solely on the close similarity between the test and pre-trained data? As I am not particularly conversant with the specific evaluation dataset utilized in this study, I am wondering how well is the transfer capability of this approach.\n2. The Finetune baseline, sharing the same mode as the Contrastive, appears to be weak, as the observed improvement over the Contrastive seems marginal according to the experimental results. I believe a more compelling comparison would be to evaluate the performance disparity between this work within the in-context learning setting versus the finetuning setting.\n3. I have concerns that the performance enhancement of this work over the baseline Contrastive may be attributed to the larger number of parameters. This work's architecture includes two message passing layers, MD and two message passing layers, MT, while the Contrastive only incorporates two message passing layers, Md. As shown in Table 4, it seems that the Contrastive model reaches saturation swiftly due to its limited capacity, while PRODIGY can persist in learning due to its greater capacity.\n", "limitations": "One limitation is that the work currently sits only on the graph classification task, and it would be nice to include some exploration of the generation task, but I won't fault the authors for that.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper explores a problem of great current interest – “how to enable in-context learning for diverse graph machine learning tasks”. \n2. The approach is novel and technically sound.  Interestingly, the \"Task graph Message Passing\" step can adjust the representation of the label nodes using the prompt examples and propagate label information back to the examples and query graph representation for test-time tasks. I think this step allows prompt to provide much stronger constraint on generation than natural language prompting.\n3. The paper is well-structured; The approach is clearly presented with descriptive figures.\n", "weaknesses": "1. The performance of the approach may depend on the similarity between the test and pre-training data, and it is not clear how well it transfers to different datasets.\n2. The fine-tuning comparison may be insufficient, as the work does not explicitly evaluate the disparity between in-context learning and fine-tuning settings for the proposed approach.\n3. The performance improvement of this work over the baseline could be partially attributed to the larger model size.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and valuable feedback. Below, we clarify a number of important points raised by the reviewer.\n\n> Re: Does the success of this approach hinge solely on the close similarity between the test and pre-trained data?\n\nThe reviewer has concerns on whether the success of our proposed framework requires close similarity between the test and pre-trained dataset. We respectfully disagree, the pre-trained data and test data are from quite diverse domains, e.g., we aim to pretrain on wikipedia knowledge graph and perform tasks on ConceptNet and NELL, both of which are drastically different from WikiKG from both node/edge feature and graph structure point of view.\n\n> Re: The Finetune baseline, sharing the same mode as the Contrastive, appears to be weak\n\nThe reviewer wonders if we adopt a weak baseline for finetune. We argue that we design our finetune baseline with contrastive learning since this is the existing best way. If we design the finetune method with our proposed method, then this is more of a variant of our own method instead of a baseline. The goal of this work is to show that with our PRODIGY framework, we for the first time allow for in-context learning over graph tasks across diverse domains, getting rid of the need of additional finetuning steps, and even outperforms traditional finetuning methods that require additional backpropagation over data on downstream domains.\n\n> Re: Number of parameters\n\nThe number of parameters for contrastive learning is ~1.8m, for our method is ~2m. We acknowledge ours has a larger number of parameters but as demonstrated in the number difference, the number of parameters of MT only accounts for a small portion of the overall model parameters. We respectfully do not think this is the main reason that leads to the performance increase.\n"}]

Author rebuttal:

Meta-review: This paper proposes a formulation of in-context learning adapted to graphs, as well as the corresponding pre-training method to enable it.

Reviews are unanimously positive, with reviewers acknowledging the novelty of the approach and the strength of the experimental results. Based on the positive feedback, I recommend acceptance.
Decision: Accept (spotlight)