SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Qian Huang\*
Hongyu Ren $^{1*}$
hyren@cs.stanford.edu
Peng Chen
pengc@stanford.edu
Gregor Krzmanc
gregor.krzmanc@ijs.si
Daniel Zeng
dzeng@cs.stanford.edu
Jure Leskovec<sup>1</sup>
jure@cs.stanford.edu
$^{2}$  University of Ljubljana
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 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 is 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 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.
In-context learning is a novel and one of the most intriguing capabilities of language models [1]. It refers to the capability of a pretrained model to 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 to solve novel tasks on novel graphs. For example, give music product
*indicates equal contribution.
Few-shot Prompting
Prompt Graph
Figure 1: In-context few-shot prompting over graphs with prompt graph for edge classification in PRODIGY. (A) Given the source graph  $\mathcal{G}$ , 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  $\mathcal{G}^{\mathrm{D}}$  by retrieving context from the source graph  $\mathcal{G}$ . (C) Then we create a task graph to capture the connection between each datapoint and each label, which includes a data node  $v_{x}$  for each datapoint and a label node  $v_{y}$  for each label in  $\mathcal{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-level 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 a 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, an in-context graph task representation, and (2) Pretraining Over Diverse In-Context 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. Prompt graph first contextualizes the input nodes/edges on which we make prediction (including both the prompt examples and the queries), then connects them with additional label nodes, such that the prompt examples are interconnected with queries. Such a unified representation allows us to specify diverse graph machine learning tasks to the same model regardless 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 to 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-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 it 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. In 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 a graph as  $\mathcal{G} = (\mathcal{V},\mathcal{E},\mathcal{R})$ , where  $\mathcal{V}$ ,  $\mathcal{E}$ ,  $\mathcal{R}$  represent the set of nodes, edges and relations. An edge  $e = (u,r,v)\in \mathcal{E}$  consists of a subject  $u\in \mathcal{V}$ , a relation  $r\in \mathcal{R}$  and an object  $v\in \mathcal{V}$ .
Given a set of classes  $\mathcal{V}$ , a standard classification task is predicting the labeling  $y\in \mathcal{V}$  of each input  $x\in \mathcal{X}$ . A node-level classification task is similar but each input is a single node in  $\mathcal{G}$ , i.e.,  $\mathcal{X} = \mathcal{V}$ , with the additional auxiliary information of the entire graph  $\mathcal{G}$ . For example, over a citation network consisting of authors and papers, a 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.,  $\mathcal{X} = \mathcal{V}\times \mathcal{V}$ . A common special case is that the classes are the same as the relations  $\mathcal{Y} = \mathcal{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  $\mathcal{G}$ .
Since we are interested in tasks of different types/levels, we design a unified formulation, where the space of the input  $\mathcal{X}$  consists of graphs, i.e.,  $x_{i} \in \mathcal{X}, x_{i} = (\mathcal{V}_{i}, \mathcal{E}_{i}, \mathcal{R}_{i})$ . For node classification task,  $\mathcal{G}_i$  only consists of the input node that we aim to make predictions on, i.e.,  $|\mathcal{V}_i| = 1$  and  $|\mathcal{E}_i| = 0$ ; for edge classification task, it consists of (subject, object) pair, i.e.,  $|\mathcal{V}_i| = 2$  and  $|\mathcal{E}_i| = 0$ .
2.2 Few-shot Prompting
Here we define in-context learning setup for classification tasks over graphs with few-shot prompting. For a  $k$ -shot prompt with a downstream  $m$ -way classification tasks with  $|\mathcal{Y}| = m$  classes, we use a small number of input-label pairs  $\mathcal{S} = \{(x_i,y_i)\}_{i=1}^{m \cdot k}$  as prompt examples of the task specification, such that there are  $k$  input-label pairs with label  $y$  for each  $y \in \mathcal{V}$ . We also give the model a set of queries  $\mathcal{Q} = \{x_i\}_{i=1}^n$  that we want to predict labels for.
We emphasize an important difference of classification tasks on graphs from language and other modalities. Namely, since all input datapoints are nodes/edges/subgraphs from the larger source graph  $\mathcal{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  $\mathcal{S}$  and  $\mathcal{Q}$ , we also need to include the source graph  $\mathcal{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  $\mathcal{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 a task graph:
Data graph. To construct a prompt graph, we first perform contextualization of each datapoint  $x_{i} = (\mathcal{V}_{i},\mathcal{E}_{i},\mathcal{R}_{i})$  in  $S$  and  $\mathcal{Q}$  in the source graph  $\mathcal{G}$  to form data graphs. The goal is to gather more information about the  $x_{i}$  from the source graph  $\mathcal{G}$  without having to represent the entire source graph explicitly. There are many potential designs for contextualization, from explicitly retrieving subgraphs to implicitly using embedding-based methods. Here we construct data graph  $\mathcal{G}_i^{\mathrm{D}}$  by sampling  $k$ -hop neighborhood of  $\mathcal{V}_i$  in  $\mathcal{G}$ . In other words,  $\mathcal{G}_i^{\mathrm{D}} = (\mathcal{V}_i^{\mathrm{D}},\mathcal{E}_i^{\mathrm{D}},\mathcal{R}_i^{\mathrm{D}})\sim \bigoplus_{i = 0}^{k}\mathsf{Neighbor}(\mathcal{V}_i,\mathcal{G},i)$ , where  $\mathcal{V}_i\subseteq \mathcal{V}_i^{\mathrm{D}}\subseteq \mathcal{V},\mathcal{E}_i\subseteq \mathcal{E}_i^{\mathrm{D}}\subseteq \mathcal{E},\mathcal{R}_i\subseteq \mathcal{R}_i^{\mathrm{D}}\subseteq \mathcal{R}$ , and Neighbor is a function that returns the exact  $i$ -hop neighbors of each node in  $\mathcal{V}_i$ . With this data graph  $\mathcal{G}_i^{\mathrm{D}}$ , we call the node set that corresponds to the nodes in  $\mathcal{V}_i$  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.
Task graph. After contextualizing each datapoint to a data graph  $\mathcal{G}^{\mathrm{D}}$ , we then construct task graph  $\mathcal{G}^{\mathrm{T}}$  to better capture the connection and relationship among the inputs and the labels. For each data graph  $\mathcal{G}_i^{\mathrm{D}}$  from the previous stage, we have a data node  $v_{x_i}$  that represents each input; for each label, we have a label node  $v_{y_i}$ . So overall, a task graph contains  $m \cdot k + n$  data nodes ( $m \cdot 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: For the query set, since we do not 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  $v_{x_i}$  will be connected to all the label nodes as shown by the yellow edges in Figure 1; 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.
Together we propose the prompt graph that consists of both data graphs and a task graph. Prompt graph effectively captures the relationship between input data  $x_{i}$  and the label  $y_{i}$  through the context captured in data graph  $\mathcal{G}_i^{\mathrm{D}}$  and the data node  $v_{x_i}$  and the label node  $v_{y_i}$  in the task graph  $\mathcal{G}^{\mathrm{T}}$ . It 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  $y_{i}$ . 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.
3 Pretraining to Enable In-context Learning
So far given a few-shot prompt for a classification task over graphs, we have defined a prompt graph representation for it that captures relationships between the prompt examples, queries, and labels. Now we need to design a pretraining strategy that can pretrain a generalizable model capable of in-context learning. We assume access to a pretraining graph  $\mathcal{G}_{\mathrm{pretrain}}$  that is independent of the source graph  $\mathcal{G}$  for the downstream task.
In this section, we introduce PRODIGY, a general pretraining framework over  $\mathcal{G}_{\mathrm{pretrain}}$  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 a message passing GNN module  $M_{\mathrm{D}}$  that learns node representation  $E$  for nodes in each  $\mathcal{G}^{\mathrm{D}}$ .
$$
E \in \mathcal {R} ^ {| \mathcal {V} ^ {\mathrm {D}} | \times d} = M _ {\mathrm {D}} (\mathcal {G} ^ {\mathrm {D}}) \tag {1}
$$
where  $d$  is the embedding dimension.  $M_{\mathrm{D}}$  can be implemented in multiple ways, such as using Graph Convolutional Network (GCN) or Graph Attention Networks (GAT) [11, 18].
To read out a single embedding  $G_{i}$  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.:
$$
G _ {i} = E _ {\mathcal {V} _ {i}} \tag {2}
$$
For link prediction tasks, we concatenate the pair of nodes, which we want to predict a link between, as well as a 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 _ {i} = W ^ {T} \left(E _ {v _ {1} \in \mathcal {V} _ {i}} \| E _ {v _ {2} \in \mathcal {V} _ {i}} \| \max  (E _ {i})\right) + b, \tag {3}
$$
where  $||$  represents concatenation,  $W\in \mathcal{R}^{3d\times d}$  is a 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  $\mathcal{G}^{\mathrm{T}}$ . We apply another GNN  $M_{\mathrm{T}}$  on the task graph to obtain updated representation of data nodes and label nodes.
$$
H = M _ {\mathrm {T}} \left(\mathcal {G} ^ {\mathrm {T}}\right) \tag {4}
$$
where  $\mathrm{H}$  is the obtained embedding per node. The initial embedding of data node  $v_{x_i}$  is  $G_i$  and the embedding of label node  $v_{y_i}$  can either be initialized with random Gaussian or additional information available about the labels. Each edge also has two binary features  $e_{ij}$  that indicate 1) whether the edge comes from an example or a query, and 2) the edge type of T or F. For  $M_{\mathrm{T}}$ , 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 a 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_{i}$  by taking cosine similarity between each pair of query graph representation and label representation, as in contrastive learning:
$$
O _ {i} = \left[ \operatorname {c o s i n e \_ s i m i l a r i t y} \left(H _ {x _ {i}}, H _ {y}\right), \forall y \in \mathcal {Y} \right] \tag {5}
$$
Note that we could perform the two message passing steps for multiple rounds to have more communication between  $x_{i}$  and learn a 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  $\mathcal{G}_{\mathrm{pretrain}}$  independent of the downstream task graph, such that the model can directly be applied on downstream tasks with in-context learning.
Our main design principle is that we formulate each pretraining objective in an in-context learning way. Most previous graph pretraining objectives only pretrain a shared graph encoder to perform various tasks with task-specific heads, so 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 to solve diverse tasks in-context with the same set of weights, such that it can perform in-context learning directly over downstream tasks.
Below, we 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 to prompt graph format (i.e. Figure 1(B,C)) with augmentation, and 3) pretraining loss over the generated prompt graph.
3.2.1 Pretraining Task Generation
We propose two methods to generate pretraining tasks from the pretraining graph  $\mathcal{G}_{\mathrm{pretrain}}$  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 a 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  $\mathcal{G}_{\mathrm{pretrain}}$  as the local neighborhoods, and we say a node belongs to a local neighborhood if it is in the sampled subgraph.
Formally, we denote  $\mathbb{N}\mathbb{M}_{k,m}$  as a sampler that generates  $m$ -way neighbor matching tasks, where each includes a  $k$ -shot prompt  $(\mathcal{G}_{\mathrm{pretrain}},\mathcal{S}_{\mathrm{NM}},\mathcal{Q}_{\mathrm{NM}})$  (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  $\mathcal{Q}_{\mathrm{NM}}$  as paired with the inputs:
$$
\left(\mathcal {G} _ {\text {p r e t r a i n}}, \mathcal {S} _ {\mathrm {N M}}, \mathcal {Q} _ {\mathrm {N M}}\right) \sim \mathrm {N M} _ {k, m} \left(\mathcal {G} _ {\text {p r e t r a i n}}\right) \tag {6}
$$
To generate these, we first sample  $m$  nodes from the pretraining graph  $\mathcal{G}_{\mathrm{pretrain}}$ , where each of the sampled node corresponds to one class.
$$
\mathcal {C} = \left\{c _ {i} \right\} _ {i = 1} ^ {m} \quad c _ {i} \sim \text {U n i f o r m} \left(\mathcal {V} _ {\text {p r e t r a i n}}\right) \tag {7}
$$
For each sampled node/class  $c_i$ , we sample  $k$  different nodes from its exact  $l$ -hop neighbors. These  $k$  nodes serve as examples of label  $c_i$ . We also sample additional  $\left\lceil \frac{n}{m} \right\rceil$  nodes similarly for each label  $c_i$  to form the query set. Formally,
$$
N _ {i} = \operatorname {N e i g h b o r} \left(c _ {i}, \mathcal {G} _ {\text {p r e t r a i n}}, l\right) \tag {8}
$$
$$
\mathcal {S} _ {i} = \left\{\left(x _ {j}, y _ {j} = c _ {i}\right) \right\} _ {j = 1} ^ {k} \quad x _ {j} \sim \text {U n i f o r m} (N _ {i}) \tag {9}
$$
$$
\mathcal {Q} _ {i} = \left\{\left(x _ {j}, y _ {j} = c _ {i}\right) \right\} _ {j = 1} ^ {\lceil \frac {n}{m} \rceil} x _ {j} \sim \text {U n i f o r m} (N _ {i}) \tag {10}
$$
In such a way, we constructed a neighbor matching pretraining task sample in the format of a few-shot prompt  $(\mathcal{G}_{\mathrm{pretrain}}, S_{\mathrm{NM}} = \bigcup S_i, Q_{\mathrm{NM}} = \bigcup Q_i)$ .
The neighbor matching task generation process outlined above is specifically applicable when the downstream tasks are also node classification. When the downstream task is link prediction, we may adapt the above neighbor matching tasks to over edges correspondingly. Specifically, we can expand each sampled input node  $x_{i}$  to an edge by randomly sampling an edge that contains  $x_{i}$ . Then, instead of classifying to 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.
Multi-task. When the pretraining graphs have node or edge-level labeling  $f(x_{i}) = y_{i} \in \mathcal{V}$  for some  $x_{i} \in \mathcal{V}_{\mathrm{pretrain}}$  or  $\mathcal{E}_{\mathrm{pretrain}}$ , 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 few-shot prompts and corresponding labels.
$$
\left(\mathcal {G} _ {\text {p r e t r a i n}}, \mathcal {S} _ {\mathrm {M T}}, \mathcal {Q} _ {\mathrm {M T}}\right) \sim \mathrm {M T} _ {k, m} \left(\mathcal {G} _ {\text {p r e t r a i n}}, f\right) \tag {11}
$$