SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For node classification tasks, we first sample  $m$  labels from the whole label set. Then, for each label, we directly sample  $k$  nodes as support examples and  $\lceil \frac{n}{m} \rceil$  nodes with labels in this set as query examples.
$$
\mathcal {C} = \left\{c _ {i} \right\} _ {i = 1} ^ {m} \quad c _ {i} \sim \text {U n i f o r m} (\mathcal {Y}) \tag {12}
$$
$$
\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} \left(\left\{x _ {i} \mid f \left(x _ {i}\right) = c _ {i} \right\}\right) \tag {13}
$$
$$
\mathcal {Q} _ {i} = \left\{\left(x _ {j}, y _ {j} = c _ {i}\right) \right\} _ {j = 1} ^ {\left[ \frac {n}{m} \right]} \quad x _ {j} \sim U n i f o r m \left(\left\{x _ {i} \mid f \left(x _ {i}\right) = c _ {i} \right\}\right) \tag {14}
$$
We then construct a task with the few-shot prompt as  $(\mathcal{G}_{\mathrm{pretrain}},\mathcal{S}_{\mathrm{MT}}) = \bigcup \mathcal{S}_i,\mathcal{Q}_{\mathrm{MT}} = \bigcup \mathcal{Q}_i)$ . For link prediction, we directly use the edge type function as  $f$ , i.e.  $f((v_{1},v_{2})) = r\Longleftrightarrow (v_{1},r,v_{2})\in \mathcal{E}$ . With this  $f$ , we may directly sample  $m$  edge types and construct pretraining tasks in a 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 as a surrogate. However, it requires extra labels if  $f$  is not part of  $\mathcal{G}_{\mathrm{pretrain}}$ , 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 (NeighborMatching and multi-task), we need to construct the prompt graph for pretraining. In addition to the standard construction process described in subsection 2.3, we add an additional augmentation step to augment
Figure 2: Neighbor matching pretraining task generation. The key idea of neighbor matching is to classify whether a node/edge is in the local neighborhood of a 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 a query node is also a two-hop neighbor of each sampled node.
Prompt Examples S: Input  $x$  Label  $y$
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  $\mathcal{G}_{\mathrm{pretrain}}$ . We first still sample the  $k$ -hop neighbor subgraph of each sample  $\mathcal{G}_i$  in the prompt examples and queries:  $\mathcal{G}_i^{\mathrm{D}} \sim \bigoplus_{j=1}^{k} \mathrm{Neighbor}(\mathcal{G}_i, \mathcal{G}_{\mathrm{pretrain}}, j)$ . Then we adopt the following two augmentation techniques to create augmented data graph  $\mathcal{G}_i^{aug}$ , 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  $\mathcal{G}_i^{aug} = \mathrm{DropNode}(\mathcal{G}_i^{\mathrm{D}})$ . For node feature masking, we randomly mask the feature of a subset of nodes with value zero to create  $\mathcal{G}_i^{aug} = \mathrm{MaskNode}(\mathcal{G}_i^{\mathrm{D}})$ . With the augmented data graphs for each datapoint in the prompt examples and the queries, we may accordingly construct the task graph  $\mathcal{G}^{\mathrm{T}}$  by creating a data node  $v_{x_i}$  for each augmented data graphs and the label node  $v_{y_i}$  as 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:
$$
\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 {15}
$$
$$
\left(\mathcal {G} _ {\text {p r e t r a i n}}, \mathcal {S} _ {\text {M T}}, \mathcal {Q} _ {\text {M T}}\right) \sim \mathrm {M T} _ {k, m} \left(\mathcal {G} _ {\text {p r e t r a i n}}, f\right) \tag {16}
$$
$$
\mathcal {L} = \underset {x _ {i} \in \mathcal {Q} _ {\mathrm {N M}}} {\mathbb {E}} \operatorname {C E} \left(O _ {\mathrm {N M}, i}, y _ {\mathrm {N M}, i}\right) + \underset {x _ {i} \in \mathcal {Q} _ {\mathrm {M T}}} {\mathbb {E}} \operatorname {C E} \left(O _ {\mathrm {M T}, i}, y _ {\mathrm {M T}, i}\right) \tag {17}
$$
where  $O_{\mathrm{NM},i}$  is the logits produced by our model over input of  $\mathcal{G}_i^{aug}$  and  $\mathcal{G}^{\mathrm{T}}$  produced from  $\mathcal{Q}_{\mathrm{NM}}$ , as described in subsection 3.1;  $y_{\mathrm{NM},i}$  is the corresponding label of  $x_i$  in  $\mathcal{Q}_{\mathrm{NM}}$ ; Similar for MT terms.
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], FB15K-237 [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-237, NELL), the downstream task is 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 different downstream datasets that belong to similar domain as the pretraining dataset (e.g., pretraining on Wiki and evaluating on ConceptNet, FB15K-237, and NELL). Each of the downstream classification datasets has its original train, validation, and test splits. To simulate the situation where there are a 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: 1) 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, we use DropNode and MaskNode augmentations with a probability of 0.5 per node for each method.
We consider three baselines for comparison: 1) NoPretrain, which uses a randomly-initialized model with the same architecture as 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
Table 1: 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 arXiv, 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 ± 0.30	65.08 ± 0.34	72.50 ± 0.35	65.64 ± 0.33	73.09 ± 0.36	65.42 ± 5.53
5	18.33 ± 0.21	51.63 ± 0.29	61.21 ± 0.28	51.97 ± 0.27	61.52 ± 0.28	53.49 ± 4.61
10	9.19 ± 0.11	36.78 ± 0.19	46.12 ± 0.19	37.23 ± 0.20	46.74 ± 0.20	30.22 ± 3.77
20	4.72 ± 0.06	25.18 ± 0.11	33.71 ± 0.12	25.91 ± 0.12	34.41 ± 0.12	17.68 ± 1.15
40	2.62 ± 0.02	17.02 ± 0.07	23.69 ± 0.06	17.19 ± 0.08	25.13 ± 0.07	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 ± 0.63	44.01 ± 0.61	46.94 ± 0.61	51.78 ± 0.63	53.97 ± 0.63	53.85 ± 9.29
5	33.54 ± 0.61	81.35 ± 0.58	80.35 ± 0.57	89.15 ± 0.46	88.02 ± 0.48	82.01 ± 12.83
10	20.0 ± 0.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.21	58.34 ± 0.22	59.58 ± 0.22	57.27 ± 3.33
5	20.95 ± 0.52	83.38 ± 0.5	82.39 ± 0.53	85.26 ± 0.48	88.09 ± 0.43	87.22 ± 12.75
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 ± 0.34	65.68 ± 0.34	68.38 ± 0.33	74.72 ± 0.31	66.19 ± 8.46
40	2.5 ± 0.06	56.7 ± 0.23	54.91 ± 0.22	51.24 ± 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 2 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 to  $48\%$  over the best baseline of Contrastive. Over KGs, PRODIGY also outperforms contrastive learning on average by  $12.2\%$ . PRODIGY also demonstrates similar-to-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 in 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 self-supervised pretraining method PG-NM produces significantly higher in-context learning performance over arXiv than baselines, even though the model is pretrained on 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 to 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 arXiv. 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\%$  over the  $65.13\%$  of TENT, even though PG-NM has never been trained on any paper category classification
Figure 3: In-context learning accuracy on ConceptNet in a 4-ways setting wrt. the number of prompt examples (shots).
Figure 4: With more training steps and data on 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).
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 all of the elements together results in 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 all ways, shown in the 'No-Attr' column.
4.4 Evaluation using different numbers of in-context examples
We investigate our method's ability to learn from the context by analyzing its performance as the number of prompt examples changes. Figure 3 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, we explore how the model scales with more pretraining data. The result on arXiv in a 5-ways setting 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.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 is called in-context learning. Comparing to previous language encoder models like BERT [4], it drastically reduces the adaptation effort comparing to fine-tuning, and has demonstrated strong performance in a 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 a 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, it then requires finetuning a 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-context 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, 9, 3, 17, 25]. However, existing meta-learning methods are only designed and tested for generalizing across different tasks on the same graph: the methods are trained on a set of training tasks on a graph, then tested over a 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.
We introduce PRODIGY, the first framework that enables in-context learning on graphs. A model that is pretrained using PRODIGY can seamlessly execute a 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.
We thank Camilo Ruiz and Michael Moor for discussions and for providing feedback on our manuscript. We also gratefully acknowledge the support of DARPA under Nos. HR00112190039 (TAMI), N660011924033 (MCS); ARO under Nos. W911NF-16-1-0342 (MURI), W911NF-16-1-0171 (DURIP); NSF under Nos. OAC-1835598 (CINES), OAC-1934578 (HDR), CCF-1918940 (Expeditions), NIH under No. 3U54HG010426-04S1 (HuBMAP), Stanford Data Science Initiative, Wu Tsai 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 is supported by Open Philanthropy AI fellowship.
The content is solely the responsibility of the authors and does not necessarily represent the official views of the funding entities.
[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 N. V. Chawla. Few-shot knowledge graph completion. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 3041-3048, 2020.
A Attribute Prediction Loss
For each augmented DataGraph  $\mathcal{G}^{aug}$ , the certain node features  $F_{v}$  are masked during MaskNode augmentation. Therefore, we can reconstruct them using the learned embedding  $E_{v}$  with a MLP and train with MSE reconstruction Loss.
$$
\mathcal {L} _ {\text {a t t r}} \left(\mathcal {G} ^ {\text {a u g}}\right) = \frac {1}{| \mathcal {V} ^ {\mathfrak {D}} |} \sum_ {v} \operatorname {M S E} \left(F _ {v}, \operatorname {M L P} \left(E _ {v}\right)\right)
$$
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
FB15K-237	15K	268K	200
NELL	69K	181K	291
C Task Graph GNN Architecture
For the GNN over task graph  $M_{\mathrm{T}}$ , we use an attention-based GNN, where each node performs attention to other nodes at each layer:
$$
\beta_ {i j} = M L P \left(W _ {q} ^ {T} H _ {i} ^ {l} \mid \mid W _ {k} ^ {T} H _ {j} ^ {l} \mid \mid e _ {i j}\right) \tag {18}
$$
$$
\alpha_ {i j} = \frac {\exp \left(\beta_ {i j}\right)}{\sum_ {k \in \mathcal {N} (i) \cup \{i \}} \exp \left(\beta_ {i k}\right)} \tag {19}
$$
$$
H _ {i} ^ {l + 1} = \operatorname {R e L U} \left(B N \left(H _ {i} ^ {l} + W _ {o} ^ {T} \sum_ {j \in \mathcal {N} (i) \cup \{i \}} \alpha_ {i j} W _ {v} ^ {T} H _ {j} ^ {l}\right)\right) \tag {20}
$$
D.1 Model Architecture, MAG240M and arxiv
We initialize the node features in citation network datasets using a 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 of two message passing layers,  $M_{\mathrm{D}}$ , over the DataGraph and one message passing layer,  $M_{\mathrm{T}}$ , over the TaskGraph. These layers are defined in Section 3.1.
For the Contrastive method, the architecture includes two message passing layers,  $M_d$ , over the DataGraph, and a 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 a linear layer head over the output of the two  $M_d$  layers, also described in Section 4.1.
D.2 Model Architecture, knowledge graph datasets
We initialize node and edge features in knowledge graph datasets using a 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,  $M_{\mathrm{D}}$ , over the DataGraph, an aggregator as described by
Table 4: Ablation of PG-NM on arXiv.
Ways	PG-NM	3→1 shot	No Attr	No Aug	No Attr, Aug	No Attr, Aug, MT
3	72.50 ± 0.35	69.13 ± 1.09	65.74 ± 1.12	68.98 ± 1.09	66.53 ± 1.12	63.60 ± 1.06
5	61.21 ± 0.29	57.49 ± 0.92	52.78 ± 0.90	57.50 ± 0.85	53.89 ± 0.92	51.27 ± 0.69
10	46.12 ± 0.19	42.03 ± 0.60	37.99 ± 0.63	42.43 ± 0.64	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 ± 0.21	21.97 ± 0.24	18.52 ± 0.22	19.69 ± 0.08
Equation 3, and two message passing layers,  $M_{\mathrm{T}}$ , over the TaskGraph, which only pass messages along the positive and query edges. These layers are defined in Section 3.1.
For the Contrastive method, the architecture includes two message passing layers,  $M_{d}$ , over the DataGraph, an aggregator as described by Equation 3 and a 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 a linear layer head over the output of the two  $M_{d}$  layers, also described in Section 4.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 to strike a balance between complexity and diversity in 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  $1 \times 10^{-3}$  and weight decay of  $1 \times 10^{-3}$ , a pretraining task with 30 ways, 3 shots, and 4 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 1 Neighbor Matching task per 1 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 of 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 3 for all evaluation tasks.
D.4 Training, Wiki
Our pretraining setup included a model with an input dimension of 768 and an embedding dimension of 256, the AdamW optimizer with a learning rate of  $1 \times 10^{-3}$  and weight decay of  $1 \times 10^{-3}$ , a pretraining task with 30 ways, 3 shots, and 4 queries per task, using a 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.